Using unbounded_type to access class object of a type annotation.

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,37 @@ 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:
+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 mypy/mypy#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.startswith(i)), 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 +1061,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/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

mypyc/test/test_function.py

@@ -0,0 +1,60 @@
+import unittest
+from unittest.mock import MagicMock, call, patch
+
+from mypy.nodes import TypeInfo
+from mypy.types import UnboundType
+from mypyc.ir.rtypes import object_rprimitive
+from mypyc.irbuild.builder import IRBuilder
+from mypyc.irbuild.function import load_type
+from mypyc.irbuild.ll_builder import LowLevelIRBuilder
+from mypyc.irbuild.mapper import Mapper
+
+
+class TestFunction(unittest.TestCase):
+    def setUp(self) -> None:
+        self.builder = MagicMock(spec=IRBuilder)
+        self.builder.configure_mock(
+            mapper=MagicMock(spec=Mapper), builder=MagicMock(spec=LowLevelIRBuilder)
+        )
+        self.builder.mapper.configure_mock(type_to_ir=[])
+        self.typ = MagicMock(spec=TypeInfo)
+        self.unbounded_type = MagicMock(spec=UnboundType)
+        self.line = 10
+
+    @patch("mypyc.irbuild.function.builtin_names", {})
+    def test_load_type_from_imported_module(self) -> None:
+        self.typ.fullname = "json.decoder.JSONDecoder"
+        self.typ.name = "JSONDecoder"
+        self.unbounded_type.name = "json.JSONDecoder"
+        self.builder.imports = {"json": "json"}
+        self.builder.load_module.return_value = "json_module"
+        self.builder.builder.get_attr.return_value = "JSONDecoder_class"
+        result = load_type(self.builder, self.typ, self.unbounded_type, self.line)
+        self.builder.load_module.assert_called_once_with("json")
+        self.builder.py_get_attr.assert_not_called()
+        self.builder.builder.get_attr.assert_called_once_with(
+            "json_module", "JSONDecoder", object_rprimitive, self.line, borrow=False
+        )
+        self.assertEqual(result, "JSONDecoder_class")
+
+    @patch("mypyc.irbuild.function.builtin_names", {})
+    def test_load_type_with_deep_nesting(self) -> None:
+        self.typ.fullname = "mod1.mod2.mod3.OuterType.InnerType"
+        self.typ.name = "InnerType"
+        self.unbounded_type.name = "mod4.mod5.mod6.OuterType.InnerType"
+        self.builder.imports = {"mod4.mod5": "mod4.mod5"}
+        self.builder.load_module.return_value = "mod4.mod5_module"
+        self.builder.py_get_attr.side_effect = ["mod4.mod5.mod6_module", "OuterType_class"]
+        self.builder.builder.get_attr.return_value = "InnerType_class"
+        result = load_type(self.builder, self.typ, self.unbounded_type, self.line)
+        self.builder.load_module.assert_called_once_with("mod4.mod5")
+        self.builder.py_get_attr.assert_has_calls(
+            [
+                call("mod4.mod5_module", "mod6", self.line),
+                call("mod4.mod5.mod6_module", "OuterType", self.line),
+            ]
+        )
+        self.builder.builder.get_attr.assert_called_once_with(
+            "OuterType_class", "InnerType", object_rprimitive, self.line, borrow=False
+        )
+        self.assertEqual(result, "InnerType_class")