Add test, add to inspect

JelleZijlstra · JelleZijlstra · commit ce98c1907a0c · 2024-05-21T21:16:45.000-04:00
diff --git a/Lib/inspect.py b/Lib/inspect.py
@@ -38,6 +38,7 @@
     "AGEN_CREATED",
     "AGEN_RUNNING",
     "AGEN_SUSPENDED",
+    "AnnotationsFormat",
     "ArgInfo",
     "Arguments",
     "Attribute",
@@ -61,6 +62,7 @@
     "ClassFoundException",
     "ClosureVars",
     "EndOfBlock",
+    "FORWARDREF",
     "FrameInfo",
     "FullArgSpec",
     "GEN_CLOSED",
@@ -134,9 +136,11 @@
     "istraceback",
     "markcoroutinefunction",
     "signature",
+    "SOURCE",
     "stack",
     "trace",
     "unwrap",
+    "VALUE",
     "walktree",
 ]
 
@@ -173,6 +177,13 @@
 TPFLAGS_IS_ABSTRACT = 1 << 20
 
 
+@enum.global_enum
+class AnnotationsFormat(enum.IntEnum):
+    VALUE = 1
+    FORWARDREF = 2
+    SOURCE = 3
+
+
 def get_annotations(obj, *, globals=None, locals=None, eval_str=False):
     """Compute the annotations dict for an object.
 
diff --git a/Lib/test/test_inspect/test_inspect.py b/Lib/test/test_inspect/test_inspect.py
@@ -1592,6 +1592,11 @@ class C(metaclass=M):
         attrs = [a[0] for a in inspect.getmembers(C)]
         self.assertNotIn('missing', attrs)
 
+    def test_annotation_format(self):
+        self.assertIs(inspect.VALUE, inspect.AnnotationsFormat.VALUE)
+        self.assertEqual(inspect.VALUE.value, 1)
+        self.assertEqual(inspect.VALUE, 1)
+
     def test_get_annotations_with_stock_annotations(self):
         def foo(a:int, b:str): pass
         self.assertEqual(inspect.get_annotations(foo), {'a': int, 'b': str})
diff --git a/Lib/test/test_type_annotations.py b/Lib/test/test_type_annotations.py
@@ -108,20 +108,24 @@ class D(metaclass=C):
         self.assertEqual(D.__annotations__, {})
 
 
+def build_module(code: str, name: str = "top") -> types.ModuleType:
+    ns = run_code(code)
+    mod = types.ModuleType(name)
+    mod.__dict__.update(ns)
+    return mod
+
+
 class TestSetupAnnotations(unittest.TestCase):
     def check(self, code: str):
         code = textwrap.dedent(code)
         for scope in ("module", "class"):
             with self.subTest(scope=scope):
                 if scope == "class":
                     code = f"class C:\n{textwrap.indent(code, '    ')}"
-                ns = run_code(code)
-                if scope == "class":
+                    ns = run_code(code)
                     annotations = ns["C"].__annotations__
                 else:
-                    mod = types.ModuleType("mod")
-                    mod.__dict__.update(ns)
-                    annotations = mod.__annotations__
+                    annotations = build_module(code).__annotations__
                 self.assertEqual(annotations, {"x": int})
 
     def test_top_level(self):
@@ -333,3 +337,20 @@ def test_no_exotic_expressions(self):
         check_syntax_error(self, "def func(x: (yield from x)): ...", "yield expression cannot be used within an annotation")
         check_syntax_error(self, "def func(x: (y := 3)): ...", "named expression cannot be used within an annotation")
         check_syntax_error(self, "def func(x: (await 42)): ...", "await expression cannot be used within an annotation")
+
+    def test_generated_annotate(self):
+        def func(x: int):
+            pass
+        class X:
+            x: int
+        mod = build_module("x: int")
+        for obj in (func, X, mod):
+            with self.subTest(obj=obj):
+                annotate = obj.__annotate__
+                self.assertIsInstance(annotate, types.FunctionType)
+                self.assertEqual(annotate.__name__, f"<annotations of {obj.__name__}>")
+                with self.assertRaises(AssertionError):  # TODO NotImplementedError
+                    annotate(2)
+                with self.assertRaises(AssertionError):  # TODO NotImplementedError
+                    annotate(None)
+                self.assertEqual(annotate(1), {"x": int})
