Track classes with explicit __module__ attributes

mypy/nodes.py

@@ -2831,6 +2831,7 @@ class is generic then it will be a type constructor of higher kind.
     __slots__ = (
         "_fullname",
         "module_name",
+        "module_override",
         "defn",
         "mro",
         "_mro_refs",
@@ -2876,6 +2877,8 @@ class is generic then it will be a type constructor of higher kind.
     # information is also in the fullname, but is harder to extract in the
     # case of nested class definitions.
     module_name: str
+    # The value of this type's __module__ attribute, if it's set explicitly.
+    module_override: str | None
     defn: ClassDef  # Corresponding ClassDef
     # Method Resolution Order: the order of looking up attributes. The first
     # value always to refers to this class.
@@ -3042,6 +3045,7 @@ def __init__(self, names: SymbolTable, defn: ClassDef, module_name: str) -> None
         self.names = names
         self.defn = defn
         self.module_name = module_name
+        self.module_override = None
         self.type_vars = []
         self.has_param_spec_type = False
         self.has_type_var_tuple_type = False
@@ -3105,6 +3109,17 @@ def name(self) -> str:
     def fullname(self) -> str:
         return self._fullname
 
+    @property
+    def qualname(self) -> str:
+        return self.fullname[len(self.module_name) + 1 :]
+
+    @property
+    def realname(self) -> str:
+        """Like fullname, but adjusted for explicit __module__ attributes"""
+        if self.module_override:
+            return f"{self.module_override}.{self.qualname}"
+        return self.fullname
+
     def is_generic(self) -> bool:
         """Is the type generic (i.e. does it have type variables)?"""
         return len(self.type_vars) > 0
@@ -3147,7 +3162,7 @@ def __getitem__(self, name: str) -> SymbolTableNode:
             raise KeyError(name)
 
     def __repr__(self) -> str:
-        return f"<TypeInfo {self.fullname}>"
+        return f"<TypeInfo {self.realname}>"
 
     def __bool__(self) -> bool:
         # We defined this here instead of just overriding it in
@@ -3253,7 +3268,7 @@ def type_str(typ: mypy.types.Type) -> str:
         if self.bases:
             base = f"Bases({', '.join(type_str(base) for base in self.bases)})"
         mro = "Mro({})".format(
-            ", ".join(item.fullname + str_conv.format_id(item) for item in self.mro)
+            ", ".join(item.realname + str_conv.format_id(item) for item in self.mro)
         )
         names = []
         for name in sorted(self.names):
@@ -3262,7 +3277,7 @@ def type_str(typ: mypy.types.Type) -> str:
             if isinstance(node, Var) and node.type:
                 description += f" ({type_str(node.type)})"
             names.append(description)
-        items = [f"Name({self.fullname})", base, mro, ("Names", names)]
+        items = [f"Name({self.realname})", base, mro, ("Names", names)]
         if self.declared_metaclass:
             items.append(f"DeclaredMetaclass({type_str(self.declared_metaclass)})")
         if self.metaclass_type:
@@ -3274,6 +3289,7 @@ def serialize(self) -> JsonDict:
         data = {
             ".class": "TypeInfo",
             "module_name": self.module_name,
+            "module_override": self.module_override,
             "fullname": self.fullname,
             "names": self.names.serialize(self.fullname),
             "defn": self.defn.serialize(),
@@ -3314,6 +3330,7 @@ def deserialize(cls, data: JsonDict) -> TypeInfo:
         module_name = data["module_name"]
         ti = TypeInfo(names, defn, module_name)
         ti._fullname = data["fullname"]
+        ti.module_override = data["module_override"]
         # TODO: Is there a reason to reconstruct ti.subtypes?
         ti.abstract_attributes = [(attr[0], attr[1]) for attr in data["abstract_attributes"]]
         ti.type_vars = data["type_vars"]

mypy/semanal.py

@@ -2893,6 +2893,7 @@ def visit_assignment_stmt(self, s: AssignmentStmt) -> None:
         self.process__all__(s)
         self.process__deletable__(s)
         self.process__slots__(s)
+        self.process__module__(s)
 
     def analyze_identity_global_assignment(self, s: AssignmentStmt) -> bool:
         """Special case 'X = X' in global scope.
@@ -4718,6 +4719,22 @@ def process__slots__(self, s: AssignmentStmt) -> None:
                 slots.extend(super_type.slots)
             self.type.slots = set(slots)
 
+    def process__module__(self, s: AssignmentStmt) -> None:
+        """
+        Processing ``__module__`` if defined in type.
+        """
+        # if isinstance(self.type, TypeInfo) and "ZoneInfo" in self.type.fullname:
+        #     breakpoint()
+        if (
+            isinstance(self.type, TypeInfo)
+            and len(s.lvalues) == 1
+            and isinstance(s.lvalues[0], NameExpr)
+            and s.lvalues[0].name == "__module__"
+            and s.lvalues[0].kind == MDEF
+            and isinstance(s.rvalue, StrExpr)
+        ):
+            self.type.module_override = s.rvalue.value
+
     #
     # Misc statements
     #

test-data/unit/semanal-typeinfo.test

@@ -89,3 +89,15 @@ TypeInfoMap(
     Mro(__main__.A, builtins.object)
     Names(
       a)))
+
+[case testExplicit__module__]
+class A:
+    __module__ = "other.module"
+[out]
+TypeInfoMap(
+  __main__.A : TypeInfo(
+    Name(other.module.A)
+    Bases(builtins.object)
+    Mro(other.module.A, builtins.object)
+    Names(
+      __module__ (builtins.str))))