Add support for __spec__ (#14739)

hauntsaninja · achimnol · web-flow · commit f5afdcd01adf · 2024-05-20T20:08:36.000-07:00
Fixes #4145 Co-authored-by: Joongi Kim <me@daybreaker.info>
diff --git a/mypy/nodes.py b/mypy/nodes.py
@@ -117,6 +117,7 @@ def set_line(
     "__file__": "__builtins__.str",
     "__package__": "__builtins__.str",
     "__annotations__": None,  # dict[str, Any] bounded in add_implicit_module_attrs()
+    "__spec__": None,  # importlib.machinery.ModuleSpec bounded in add_implicit_module_attrs()
 }
 
 
diff --git a/mypy/semanal.py b/mypy/semanal.py
@@ -283,6 +283,7 @@
     TypeVarTupleType,
     TypeVarType,
     UnboundType,
+    UnionType,
     UnpackType,
     get_proper_type,
     get_proper_types,
@@ -635,6 +636,7 @@ def add_implicit_module_attrs(self, file_node: MypyFile) -> None:
         str_type: Type | None = self.named_type_or_none("builtins.str")
         if str_type is None:
             str_type = UnboundType("builtins.str")
+        inst: Type | None
         for name, t in implicit_module_attrs.items():
             if name == "__doc__":
                 typ: Type = str_type
@@ -660,6 +662,22 @@ def add_implicit_module_attrs(self, file_node: MypyFile) -> None:
                     self.defer()
                     return
                 typ = inst
+            elif name == "__spec__":
+                if self.options.use_builtins_fixtures:
+                    inst = self.named_type_or_none("builtins.object")
+                else:
+                    inst = self.named_type_or_none("importlib.machinery.ModuleSpec")
+                if inst is None:
+                    if self.final_iteration:
+                        inst = self.named_type_or_none("builtins.object")
+                        assert inst is not None, "Cannot find builtins.object"
+                    else:
+                        self.defer()
+                        return
+                if file_node.name == "__main__":
+                    # https://docs.python.org/3/reference/import.html#main-spec
+                    inst = UnionType.make_union([inst, NoneType()])
+                typ = inst
             else:
                 assert t is not None, f"type should be specified for {name}"
                 typ = UnboundType(t)
diff --git a/test-data/unit/check-basic.test b/test-data/unit/check-basic.test
@@ -219,6 +219,14 @@ reveal_type(__doc__)          # N: Revealed type is "builtins.str"
 reveal_type(__file__)         # N: Revealed type is "builtins.str"
 reveal_type(__package__)      # N: Revealed type is "builtins.str"
 reveal_type(__annotations__)  # N: Revealed type is "builtins.dict[builtins.str, Any]"
+# This will actually reveal Union[importlib.machinery.ModuleSpec, None]
+reveal_type(__spec__)         # N: Revealed type is "Union[builtins.object, None]"
+
+import module
+reveal_type(module.__name__)  # N: Revealed type is "builtins.str"
+# This will actually reveal importlib.machinery.ModuleSpec
+reveal_type(module.__spec__)  # N: Revealed type is "builtins.object"
+[file module.py]
 [builtins fixtures/primitives.pyi]
 
 
diff --git a/test-data/unit/fine-grained-inspect.test b/test-data/unit/fine-grained-inspect.test
@@ -236,7 +236,7 @@ class C: ...
 [builtins fixtures/module.pyi]
 [out]
 ==
-{"<pack.bar>": ["C", "__annotations__", "__doc__", "__file__", "__name__", "__package__", "bar", "x"], "ModuleType": ["__file__", "__getattr__"]}
+{"<pack.bar>": ["C", "__annotations__", "__doc__", "__file__", "__name__", "__package__", "__spec__", "bar", "x"], "ModuleType": ["__file__", "__getattr__"]}
 
 [case testInspectModuleDef]
 # inspect2: --show=definition --include-kind tmp/foo.py:2:1
diff --git a/test-data/unit/pythoneval.test b/test-data/unit/pythoneval.test
@@ -107,14 +107,20 @@ f
 [case testModuleAttributes]
 import math
 import typing
+print(type(__spec__))
 print(math.__name__)
+print(math.__spec__.name)
 print(type(math.__dict__))
 print(type(math.__doc__ or ''))
+print(type(math.__spec__).__name__)
 print(math.__class__)
 [out]
+<class 'NoneType'>
+math
 math
 <class 'dict'>
 <class 'str'>
+ModuleSpec
 <class 'module'>
 
 [case testSpecialAttributes]

Original file line number	Diff line number	Diff line change
`@@ -117,6 +117,7 @@ def set_line(`
`117`	`117`	`"__file__": "__builtins__.str",`
`118`	`118`	`"__package__": "__builtins__.str",`
`119`	`119`	`"__annotations__": None, # dict[str, Any] bounded in add_implicit_module_attrs()`
	`120`	`+ "__spec__": None, # importlib.machinery.ModuleSpec bounded in add_implicit_module_attrs()`
`120`	`121`	`}`
`121`	`122`
`122`	`123`