diff --git a/ChangeLog b/ChangeLog index 8ddaae2ae7..46d5a91a5e 100644 --- a/ChangeLog +++ b/ChangeLog @@ -12,6 +12,11 @@ What's New in astroid 2.9.1? ============================ Release date: TBA +* Treat ``typing.NewType()`` values as normal subclasses. + + Closes PyCQA/pylint#2296 + Closes PyCQA/pylint#3162 + * Prefer the module loader get_source() method in AstroidBuilder's module_build() when possible to avoid assumptions about source code being available on a filesystem. Otherwise the source cannot diff --git a/astroid/brain/brain_typing.py b/astroid/brain/brain_typing.py index 2c4d031763..60079b56ef 100644 --- a/astroid/brain/brain_typing.py +++ b/astroid/brain/brain_typing.py @@ -38,8 +38,6 @@ from astroid.util import Uninferable TYPING_NAMEDTUPLE_BASENAMES = {"NamedTuple", "typing.NamedTuple"} -TYPING_TYPEVARS = {"TypeVar", "NewType"} -TYPING_TYPEVARS_QUALIFIED = {"typing.TypeVar", "typing.NewType"} TYPING_TYPE_TEMPLATE = """ class Meta(type): def __getitem__(self, item): @@ -52,6 +50,11 @@ def __args__(self): class {0}(metaclass=Meta): pass """ +# PEP484 suggests NewType is equivalent to this for typing purposes +TYPING_NEWTYPE_TEMPLATE = """ +class {derived}({base}): + pass +""" TYPING_MEMBERS = set(getattr(typing, "__all__", [])) TYPING_ALIAS = frozenset( @@ -106,23 +109,32 @@ def __class_getitem__(cls, item): """ -def looks_like_typing_typevar_or_newtype(node): +def looks_like_typing_typevar(node): + func = node.func + if isinstance(func, Attribute): + return func.attrname == "TypeVar" + if isinstance(func, Name): + return func.name == "TypeVar" + return False + + +def looks_like_typing_newtype(node): func = node.func if isinstance(func, Attribute): - return func.attrname in TYPING_TYPEVARS + return func.attrname == "NewType" if isinstance(func, Name): - return func.name in TYPING_TYPEVARS + return func.name == "NewType" return False -def infer_typing_typevar_or_newtype(node, context_itton=None): +def infer_typing_typevar(node, context_itton=None): """Infer a typing.TypeVar(...) or typing.NewType(...) call""" try: func = next(node.func.infer(context=context_itton)) except (InferenceError, StopIteration) as exc: raise UseInferenceDefault from exc - if func.qname() not in TYPING_TYPEVARS_QUALIFIED: + if func.qname() != "typing.TypeVar": raise UseInferenceDefault if not node.args: raise UseInferenceDefault @@ -132,6 +144,24 @@ def infer_typing_typevar_or_newtype(node, context_itton=None): return node.infer(context=context_itton) +def infer_typing_newtype(node, context_itton=None): + """Infer a typing.TypeVar(...) or typing.NewType(...) call""" + try: + func = next(node.func.infer(context=context_itton)) + except (InferenceError, StopIteration) as exc: + raise UseInferenceDefault from exc + + if func.qname() != "typing.NewType": + raise UseInferenceDefault + if not node.args: + raise UseInferenceDefault + + derived = node.args[0].as_string().strip("'") + base = node.args[1].as_string().strip("'") + node = extract_node(TYPING_NEWTYPE_TEMPLATE.format(derived=derived, base=base)) + return node.infer(context=context_itton) + + def _looks_like_typing_subscript(node): """Try to figure out if a Subscript node *might* be a typing-related subscript""" if isinstance(node, Name): @@ -409,8 +439,13 @@ def infer_typing_cast( AstroidManager().register_transform( Call, - inference_tip(infer_typing_typevar_or_newtype), - looks_like_typing_typevar_or_newtype, + inference_tip(infer_typing_typevar), + looks_like_typing_typevar, +) +AstroidManager().register_transform( + Call, + inference_tip(infer_typing_newtype), + looks_like_typing_newtype, ) AstroidManager().register_transform( Subscript, inference_tip(infer_typing_attr), _looks_like_typing_subscript diff --git a/tests/unittest_brain.py b/tests/unittest_brain.py index d38c0d79b4..e3bda24aa1 100644 --- a/tests/unittest_brain.py +++ b/tests/unittest_brain.py @@ -1659,6 +1659,34 @@ def test_typing_types(self) -> None: inferred = next(node.infer()) self.assertIsInstance(inferred, nodes.ClassDef, node.as_string()) + def test_typing_newtype_attrs(self) -> None: + ast_nodes = builder.extract_node( + """ + from typing import List, NewType + import typing + NewType("Foo", str) #@ + NewType("Bar", "int") #@ + """ + ) + assert isinstance(ast_nodes, list) + + # Base type given by reference + foo_node = ast_nodes[0] + foo_inferred = next(foo_node.infer()) + self.assertIsInstance(foo_inferred, astroid.ClassDef) + + # Check base type method is inferred + foo_base_class_method = foo_inferred.getattr("endswith")[0] + self.assertIsInstance(foo_base_class_method, astroid.FunctionDef) + + # Base type given by string + bar_node = ast_nodes[1] + bar_inferred = next(bar_node.infer()) + self.assertIsInstance(bar_inferred, astroid.ClassDef) + + bar_base_class_method = bar_inferred.getattr("bit_count")[0] + self.assertIsInstance(bar_base_class_method, astroid.FunctionDef) + def test_namedtuple_nested_class(self): result = builder.extract_node( """