Skip to content

Commit

Permalink
Track classes with explicit __module__ attributes
Browse files Browse the repository at this point in the history
Having access to this information can be used for more accurate
diagnostics and tests
  • Loading branch information
tungol committed Dec 21, 2023
1 parent 7d842e8 commit d0ec2c4
Show file tree
Hide file tree
Showing 3 changed files with 49 additions and 3 deletions.
23 changes: 20 additions & 3 deletions mypy/nodes.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down Expand Up @@ -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
# This is set if the class contains an explicit __module__ attribute
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.
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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):
Expand All @@ -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:
Expand All @@ -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(),
Expand Down Expand Up @@ -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"]
Expand Down
17 changes: 17 additions & 0 deletions mypy/semanal.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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
#
Expand Down
12 changes: 12 additions & 0 deletions test-data/unit/semanal-typeinfo.test
Original file line number Diff line number Diff line change
Expand Up @@ -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))))

0 comments on commit d0ec2c4

Please sign in to comment.