Improve NewType inference for string forward refs

Author: colatkinson

astroid/brain/brain_typing.py

@@ -14,6 +14,7 @@
 from astroid.builder import _extract_single_node
 from astroid.const import PY38_PLUS, PY39_PLUS
 from astroid.exceptions import (
+    AstroidImportError,
     AttributeInferenceError,
     InferenceError,
     UseInferenceDefault,
@@ -171,27 +172,109 @@ def infer_typing_newtype(
     )
     new_node.parent = node.parent
 
-    # Base type arg is a normal reference, so no need to do special lookups
-    if not isinstance(base, nodes.Const):
-        new_node.postinit(
-            bases=[base], body=new_node.body, decorators=new_node.decorators
-        )
+    new_bases: list[NodeNG] = []
 
-    # If the base type is given as a string (e.g. for a forward reference),
-    # make a naive attempt to find the corresponding node.
-    # Note that this will not work with imported types.
-    if isinstance(base, nodes.Const) and isinstance(base.value, str):
+    if not isinstance(base, nodes.Const):
+        # Base type arg is a normal reference, so no need to do special lookups
+        new_bases = [base]
+    elif isinstance(base, nodes.Const) and isinstance(base.value, str):
+        # If the base type is given as a string (e.g. for a forward reference),
+        # make a naive attempt to find the corresponding node.
         _, resolved_base = node.frame().lookup(base_name)
         if resolved_base:
-            new_node.postinit(
-                bases=[resolved_base[0]],
-                body=new_node.body,
-                decorators=new_node.decorators,
-            )
+            base_node = resolved_base[0]
+
+            # If the value is from an "import from" statement, follow the import chain
+            if isinstance(base_node, nodes.ImportFrom):
+                ctx = ctx.clone() if ctx else context.InferenceContext()
+                ctx.lookupname = base_name
+                base_node = next(base_node.infer(context=ctx))
+
+            new_bases = [base_node]
+        elif "." in base.value:
+            possible_base = _try_find_imported_object_from_str(node, base.value, ctx)
+            if possible_base:
+                new_bases = [possible_base]
+
+    if new_bases:
+        new_node.postinit(
+            bases=new_bases, body=new_node.body, decorators=new_node.decorators
+        )
 
     return new_node.infer(context=ctx)
 
 
+def _try_find_imported_object_from_str(
+    node: nodes.Call,
+    name: str,
+    ctx: context.InferenceContext | None,
+) -> nodes.NodeNG | None:
+    for statement_mod_name, _ in _possible_module_object_splits(name):
+        # Find import statements that may pull in the appropriate modules
+        # The name used to find this statement may not correspond to the name of the module actually being imported
+        # For example, "import email.charset" is found by lookup("email")
+        _, resolved_bases = node.frame().lookup(statement_mod_name)
+        if not resolved_bases:
+            continue
+
+        resolved_base = resolved_bases[0]
+        if isinstance(resolved_base, nodes.Import):
+            # Extract the names of the module as they are accessed from actual code
+            scope_names = {(alias or name) for (name, alias) in resolved_base.names}
+            aliases = {alias: name for (name, alias) in resolved_base.names if alias}
+
+            # Find potential mod_name, obj_name splits that work with the available names
+            # for the module in this scope
+            import_targets = [
+                (mod_name, obj_name)
+                for (mod_name, obj_name) in _possible_module_object_splits(name)
+                if mod_name in scope_names
+            ]
+            if not import_targets:
+                continue
+
+            import_target, name_in_mod = import_targets[0]
+            import_target = aliases.get(import_target, import_target)
+
+            # Try to import the module and find the object in it
+            try:
+                resolved_mod: nodes.Module = resolved_base.do_import_module(
+                    import_target
+                )
+            except AstroidImportError:
+                # If the module doesn't actually exist, try the next option
+                continue
+
+            # Try to find the appropriate ClassDef or other such node in the target module
+            _, object_results_in_mod = resolved_mod.lookup(name_in_mod)
+            if not object_results_in_mod:
+                continue
+
+            base_node = object_results_in_mod[0]
+
+            # If the value is from an "import from" statement, follow the import chain
+            if isinstance(base_node, nodes.ImportFrom):
+                ctx = ctx.clone() if ctx else context.InferenceContext()
+                ctx.lookupname = name_in_mod
+                base_node = next(base_node.infer(context=ctx))
+
+            return base_node
+
+    return None
+
+
+def _possible_module_object_splits(
+    dot_str: str,
+) -> Iterator[tuple[str, str]]:
+    components = dot_str.split(".")
+    popped = []
+
+    while components:
+        popped.append(components.pop())
+
+        yield ".".join(components), ".".join(reversed(popped))
+
+
 def _looks_like_typing_subscript(node):
     """Try to figure out if a Subscript node *might* be a typing-related subscript"""
     if isinstance(node, Name):

tests/unittest_brain.py

@@ -1837,6 +1837,172 @@ def _verify_node_has_expected_attr(self, node: nodes.NodeNG) -> None:
         with self.assertRaises(AttributeInferenceError):
             inferred.getattr("bad_attr")
 
+    def test_typing_newtype_forward_reference_imported(self) -> None:
+        all_ast_nodes = builder.extract_node(
+            """
+        from typing import NewType
+
+        A = NewType("A", "decimal.Decimal")
+        B = NewType("B", "decimal_mod_alias.Decimal")
+        C = NewType("C", "Decimal")
+        D = NewType("D", "DecimalAlias")
+
+        import decimal
+        import decimal as decimal_mod_alias
+        from decimal import Decimal
+        from decimal import Decimal as DecimalAlias
+
+        Decimal #@
+
+        a = A(decimal.Decimal(2))
+        a #@
+        b = B(decimal_mod_alias.Decimal(2))
+        b #@
+        c = C(Decimal(2))
+        c #@
+        d = D(DecimalAlias(2))
+        d #@
+        """
+        )
+        assert isinstance(all_ast_nodes, list)
+
+        real_dec, *ast_nodes = all_ast_nodes
+
+        real_quantize = next(real_dec.infer()).getattr("quantize")
+
+        for node in ast_nodes:
+            all_inferred = list(node.infer())
+            assert len(all_inferred) == 1
+            inferred = all_inferred[0]
+            assert isinstance(inferred, astroid.Instance)
+
+            assert inferred.getattr("quantize") == real_quantize
+
+    def test_typing_newtype_forward_ref_bad_base(self) -> None:
+        ast_nodes = builder.extract_node(
+            """
+        from typing import NewType
+
+        A = NewType("A", "DoesntExist")
+
+        a = A()
+        a #@
+
+        # Valid name, but not actually imported
+        B = NewType("B", "decimal.Decimal")
+
+        b = B()
+        b #@
+
+        # AST works out, but can't import the module
+        import not_a_real_module
+
+        C = NewType("C", "not_a_real_module.SomeClass")
+        c = C()
+        c #@
+
+        # Real module, fake base class name
+        import email.charset
+
+        D = NewType("D", "email.charset.BadClassRef")
+        d = D()
+        d #@
+
+        # Real module, but aliased differently than used
+        import email.header as header_mod
+
+        E = NewType("E", "email.header.Header")
+        e = E(header_mod.Header())
+        e #@
+        """
+        )
+        assert isinstance(ast_nodes, list)
+
+        for ast_node in ast_nodes:
+            inferred = next(ast_node.infer())
+
+            with self.assertRaises(astroid.AttributeInferenceError):
+                inferred.getattr("value")
+
+    def test_typing_newtype_forward_ref_nested_module(self) -> None:
+        ast_nodes = builder.extract_node(
+            """
+        from typing import NewType
+
+        A = NewType("A", "email.charset.Charset")
+        B = NewType("B", "charset.Charset")
+
+        # header is unused in both cases, but verifies that module name is properly checked
+        import email.header, email.charset
+        from email import header, charset
+
+        real = charset.Charset()
+        real #@
+
+        a = A(email.charset.Charset())
+        a #@
+
+        b = B(charset.Charset())
+        """
+        )
+        assert isinstance(ast_nodes, list)
+
+        real, *newtypes = ast_nodes
+
+        real_inferred_all = list(real.infer())
+        assert len(real_inferred_all) == 1
+        real_inferred = real_inferred_all[0]
+
+        real_method = real_inferred.getattr("get_body_encoding")
+
+        for newtype_node in newtypes:
+            newtype_inferred_all = list(newtype_node.infer())
+            assert len(newtype_inferred_all) == 1
+            newtype_inferred = newtype_inferred_all[0]
+
+            newtype_method = newtype_inferred.getattr("get_body_encoding")
+
+            assert real_method == newtype_method
+
+    def test_typing_newtype_forward_ref_nested_class(self) -> None:
+        ast_nodes = builder.extract_node(
+            """
+        from typing import NewType
+
+        A = NewType("A", "SomeClass.Nested")
+
+        class SomeClass:
+            class Nested:
+                def method(self) -> None:
+                    pass
+
+        real = SomeClass.Nested()
+        real #@
+
+        a = A(SomeClass.Nested())
+        a #@
+        """
+        )
+        assert isinstance(ast_nodes, list)
+
+        real, newtype = ast_nodes
+
+        real_all_inferred = list(real.infer())
+        assert len(real_all_inferred) == 1
+        real_inferred = real_all_inferred[0]
+        real_method = real_inferred.getattr("method")
+
+        newtype_all_inferred = list(newtype.infer())
+        assert len(newtype_all_inferred) == 1
+        newtype_inferred = newtype_all_inferred[0]
+
+        # This could theoretically work, but for now just here to check that
+        # the "forward-declared module" inference doesn't totally break things
+        with self.assertRaises(astroid.AttributeInferenceError):
+            newtype_method = newtype_inferred.getattr("method")
+
+            assert real_method == newtype_method
+
     def test_namedtuple_nested_class(self):
         result = builder.extract_node(
             """