[mypyc] Fix unsupported imports for type annotations (#20390)

This is a follow up to #18874 which broke certain imports of type
annotations.

Fixes mypyc/mypyc#1099
Fixes mypyc/mypyc#1167
Fixes #20365
Closes mypyc/mypyc#1100

Author: lgeiger

mypyc/irbuild/function.py

@@ -858,11 +858,6 @@ def get_func_target(builder: IRBuilder, fdef: FuncDef) -> AssignmentTarget:
     return builder.add_local_reg(fdef, object_rprimitive)
 
 
-# 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
@@ -876,34 +871,17 @@ def load_type(builder: IRBuilder, typ: TypeInfo, unbounded_type: Type | None, li
     # `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))
-    # 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
-        )
+    elif isinstance(unbounded_type, UnboundType):
+        path_parts = unbounded_type.name.split(".")
+        class_obj = builder.load_global_str(path_parts[0], line)
+        for attr in path_parts[1:]:
+            class_obj = builder.py_get_attr(class_obj, attr, line)
     else:
         class_obj = builder.load_global_str(typ.name, line)
 

mypyc/test-data/run-classes.test

@@ -95,6 +95,36 @@ print(C.__annotations__["spam"] is JSONDecoder)
 True
 True
 
+[case testTypedDictWithFieldsAs]
+import json as _json
+from typing import TypedDict
+
+class C(TypedDict):
+    spam: _json.JSONDecoder
+[file driver.py]
+from native import C
+from json import JSONDecoder
+
+print(C.__annotations__["spam"] is JSONDecoder)
+[typing fixtures/typing-full.pyi]
+[out]
+True
+
+[case testTypedDictWithFieldsFrom]
+from json import decoder
+from typing import TypedDict
+
+class C(TypedDict):
+    spam: decoder.JSONDecoder
+[file driver.py]
+from native import C
+from json import JSONDecoder
+
+print(C.__annotations__["spam"] is JSONDecoder)
+[typing fixtures/typing-full.pyi]
+[out]
+True
+
 [case testClassWithDeletableAttributes]
 from typing import Any, cast
 from testutil import assertRaises