[mypyc] Using UnboundedType to access class object of a type annotation. (python#18874)

Fixes mypyc/mypyc#1087.

This fix handles cases where type annotation is nested inside an
imported module (like an inner class) or imported from a different
module.

Author: advait-dixit

mypyc/irbuild/classdef.py

@@ -634,7 +634,7 @@ def add_non_ext_class_attr_ann(
             if builder.current_module == type_info.module_name and stmt.line < type_info.line:
                 typ = builder.load_str(type_info.fullname)
             else:
-                typ = load_type(builder, type_info, stmt.line)
+                typ = load_type(builder, type_info, stmt.unanalyzed_type, stmt.line)
 
     if typ is None:
         # FIXME: if get_type_info is not provided, don't fall back to stmt.type?
@@ -650,7 +650,7 @@ def add_non_ext_class_attr_ann(
             # actually a forward reference due to the __annotations__ future?
             typ = builder.load_str(stmt.unanalyzed_type.original_str_expr)
         elif isinstance(ann_type, Instance):
-            typ = load_type(builder, ann_type.type, stmt.line)
+            typ = load_type(builder, ann_type.type, stmt.unanalyzed_type, stmt.line)
         else:
             typ = builder.add(LoadAddress(type_object_op.type, type_object_op.src, stmt.line))
 

mypyc/irbuild/function.py

@@ -29,7 +29,7 @@
     TypeInfo,
     Var,
 )
-from mypy.types import CallableType, get_proper_type
+from mypy.types import CallableType, Type, UnboundType, get_proper_type
 from mypyc.common import LAMBDA_NAME, PROPSET_PREFIX, SELF_NAME
 from mypyc.ir.class_ir import ClassIR, NonExtClassInfo
 from mypyc.ir.func_ir import (
@@ -802,15 +802,49 @@ def get_func_target(builder: IRBuilder, fdef: FuncDef) -> AssignmentTarget:
     return builder.add_local_reg(fdef, object_rprimitive)
 
 
-def load_type(builder: IRBuilder, typ: TypeInfo, line: int) -> Value:
+# This function still does not support the following imports.
+# import json as _json
+# from json import decoder
+# Using either _json.JSONDecoder or decoder.JSONDecoder as a type hint for a dataclass field will fail.
+# See issue mypyc/mypyc#1099.
+def load_type(builder: IRBuilder, typ: TypeInfo, unbounded_type: Type | None, line: int) -> Value:
+    # typ.fullname contains the module where the class object was defined. However, it is possible
+    # that the class object's module was not imported in the file currently being compiled. So, we
+    # use unbounded_type.name (if provided by caller) to load the class object through one of the
+    # imported modules.
+    # Example: for `json.JSONDecoder`, typ.fullname is `json.decoder.JSONDecoder` but the Python
+    # file may import `json` not `json.decoder`.
+    # Another corner case: The Python file being compiled imports mod1 and has a type hint
+    # `mod1.OuterClass.InnerClass`. But, mod1/__init__.py might import OuterClass like this:
+    # `from mod2.mod3 import OuterClass`. In this case, typ.fullname is
+    # `mod2.mod3.OuterClass.InnerClass` and `unbounded_type.name` is `mod1.OuterClass.InnerClass`.
+    # So, we must use unbounded_type.name to load the class object.
+    # See issue mypyc/mypyc#1087.
+    load_attr_path = (
+        unbounded_type.name if isinstance(unbounded_type, UnboundType) else typ.fullname
+    ).removesuffix(f".{typ.name}")
     if typ in builder.mapper.type_to_ir:
         class_ir = builder.mapper.type_to_ir[typ]
         class_obj = builder.builder.get_native_type(class_ir)
     elif typ.fullname in builtin_names:
         builtin_addr_type, src = builtin_names[typ.fullname]
         class_obj = builder.add(LoadAddress(builtin_addr_type, src, line))
-    elif typ.module_name in builder.imports:
-        loaded_module = builder.load_module(typ.module_name)
+    # This elif-condition finds the longest import that matches the load_attr_path.
+    elif module_name := max(
+        (i for i in builder.imports if load_attr_path == i or load_attr_path.startswith(f"{i}.")),
+        default="",
+        key=len,
+    ):
+        # Load the imported module.
+        loaded_module = builder.load_module(module_name)
+        # Recursively load attributes of the imported module. These may be submodules, classes or
+        # any other object.
+        for attr in (
+            load_attr_path.removeprefix(f"{module_name}.").split(".")
+            if load_attr_path != module_name
+            else []
+        ):
+            loaded_module = builder.py_get_attr(loaded_module, attr, line)
         class_obj = builder.builder.get_attr(
             loaded_module, typ.name, object_rprimitive, line, borrow=False
         )
@@ -1039,7 +1073,7 @@ def maybe_insert_into_registry_dict(builder: IRBuilder, fitem: FuncDef) -> None:
         )
         registry = load_singledispatch_registry(builder, dispatch_func_obj, line)
         for typ in types:
-            loaded_type = load_type(builder, typ, line)
+            loaded_type = load_type(builder, typ, None, line)
             builder.primitive_op(dict_set_item_op, [registry, loaded_type, to_insert], line)
         dispatch_cache = builder.builder.get_attr(
             dispatch_func_obj, "dispatch_cache", dict_rprimitive, line

mypyc/test-data/commandline.test

@@ -261,3 +261,31 @@ print("imported foo")
 importing...
 imported foo
 done
+
+[case testImportFromInitPy]
+# cmd: foo.py
+import foo
+
+[file pkg2/__init__.py]
+
+[file pkg2/mod2.py]
+class A:
+    class B:
+        pass
+
+[file pkg1/__init__.py]
+from pkg2.mod2 import A
+
+[file foo.py]
+import pkg1
+from typing import TypedDict
+
+class Eggs(TypedDict):
+    obj1: pkg1.A.B
+
+print(type(Eggs(obj1=pkg1.A.B())["obj1"]).__name__)
+print(type(Eggs(obj1=pkg1.A.B())["obj1"]).__module__)
+
+[out]
+B
+pkg2.mod2

mypyc/test-data/run-classes.test

@@ -78,17 +78,22 @@ assert hasattr(c, 'x')
 
 [case testTypedDictWithFields]
 import collections
+import json
 from typing import TypedDict
 class C(TypedDict):
     x: collections.deque
+    spam: json.JSONDecoder
 [file driver.py]
 from native import C
 from collections import deque
+from json import JSONDecoder
 
 print(C.__annotations__["x"] is deque)
+print(C.__annotations__["spam"] is JSONDecoder)
 [typing fixtures/typing-full.pyi]
 [out]
 True
+True
 
 [case testClassWithDeletableAttributes]
 from typing import Any, cast