Improve NewType inference for string forward refs

colatkinson · colatkinson · commit a1370de4ff64 · 2022-03-07T06:31:49.000-05:00
diff --git a/astroid/brain/brain_typing.py b/astroid/brain/brain_typing.py
@@ -22,6 +22,7 @@
 from astroid.builder import _extract_single_node
 from astroid.const import PY37_PLUS, PY38_PLUS, PY39_PLUS
 from astroid.exceptions import (
+    AstroidImportError,
     AttributeInferenceError,
     InferenceError,
     UseInferenceDefault,
@@ -191,19 +192,115 @@ def infer_typing_newtype(
 
     # 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):
         _, resolved_base = node.frame().lookup(base_name)
         if resolved_base:
+            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 = (
+                    context_itton.clone()
+                    if context_itton
+                    else context.InferenceContext()
+                )
+                ctx.lookupname = base_name
+                base_node = next(base_node.infer(context=ctx))
+
             new_node.postinit(
-                bases=[resolved_base[0]],
+                bases=[base_node],
                 body=new_node.body,
                 decorators=new_node.decorators,
             )
+        elif "." in base.value:
+            possible_base = _try_find_imported_object_from_str(
+                node, base.value, context_itton
+            )
+            if possible_base:
+                new_node.postinit(
+                    bases=[possible_base],
+                    body=new_node.body,
+                    decorators=new_node.decorators,
+                )
 
     return new_node.infer(context=context_itton)
 
 
+def _try_find_imported_object_from_str(
+    node: nodes.Call,
+    name: str,
+    context_itton: typing.Optional[context.InferenceContext],
+) -> typing.Optional[nodes.NodeNG]:
+    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 = (
+                    context_itton.clone()
+                    if context_itton
+                    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,
+) -> typing.Iterator[typing.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):
diff --git a/tests/unittest_brain.py b/tests/unittest_brain.py
@@ -1817,6 +1817,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]
+            self.assertIsInstance(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(
             """
