gh-125618: Make FORWARDREF format succeed more often (#132818)

Fixes #125618.

1 files changed, 92 insertions, 8 deletions

diff --git a/Lib/test/test_annotationlib.py b/Lib/test/test_annotationlib.py
index d9000b63922..13c6a2a584b 100644
--- a/Lib/test/test_annotationlib.py
+++ b/Lib/test/test_annotationlib.py
@@ -276,10 +276,10 @@ class TestStringFormat(unittest.TestCase):
     def test_getitem(self):
         def f(x: undef1[str, undef2]):
             pass
-        anno = annotationlib.get_annotations(f, format=Format.STRING)
+        anno = get_annotations(f, format=Format.STRING)
         self.assertEqual(anno, {"x": "undef1[str, undef2]"})
 
-        anno = annotationlib.get_annotations(f, format=Format.FORWARDREF)
+        anno = get_annotations(f, format=Format.FORWARDREF)
         fwdref = anno["x"]
         self.assertIsInstance(fwdref, ForwardRef)
         self.assertEqual(
@@ -289,18 +289,18 @@ class TestStringFormat(unittest.TestCase):
     def test_slice(self):
         def f(x: a[b:c]):
             pass
-        anno = annotationlib.get_annotations(f, format=Format.STRING)
+        anno = get_annotations(f, format=Format.STRING)
         self.assertEqual(anno, {"x": "a[b:c]"})
 
         def f(x: a[b:c, d:e]):
             pass
-        anno = annotationlib.get_annotations(f, format=Format.STRING)
+        anno = get_annotations(f, format=Format.STRING)
         self.assertEqual(anno, {"x": "a[b:c, d:e]"})
 
         obj = slice(1, 1, 1)
         def f(x: obj):
             pass
-        anno = annotationlib.get_annotations(f, format=Format.STRING)
+        anno = get_annotations(f, format=Format.STRING)
         self.assertEqual(anno, {"x": "obj"})
 
     def test_literals(self):
@@ -316,7 +316,7 @@ class TestStringFormat(unittest.TestCase):
         ):
             pass
 
-        anno = annotationlib.get_annotations(f, format=Format.STRING)
+        anno = get_annotations(f, format=Format.STRING)
         self.assertEqual(
             anno,
             {
@@ -335,7 +335,7 @@ class TestStringFormat(unittest.TestCase):
         # Simple case first
         def f(x: a[[int, str], float]):
             pass
-        anno = annotationlib.get_annotations(f, format=Format.STRING)
+        anno = get_annotations(f, format=Format.STRING)
         self.assertEqual(anno, {"x": "a[[int, str], float]"})
 
         def g(
@@ -345,7 +345,7 @@ class TestStringFormat(unittest.TestCase):
             z: a[(int, str), 5],
         ):
             pass
-        anno = annotationlib.get_annotations(g, format=Format.STRING)
+        anno = get_annotations(g, format=Format.STRING)
         self.assertEqual(
             anno,
             {
@@ -1017,6 +1017,58 @@ class TestGetAnnotations(unittest.TestCase):
             set(results.generic_func.__type_params__),
         )
 
+    def test_partial_evaluation(self):
+        def f(
+            x: builtins.undef,
+            y: list[int],
+            z: 1 + int,
+            a: builtins.int,
+            b: [builtins.undef, builtins.int],
+        ):
+            pass
+
+        self.assertEqual(
+            get_annotations(f, format=Format.FORWARDREF),
+            {
+                "x": support.EqualToForwardRef("builtins.undef", owner=f),
+                "y": list[int],
+                "z": support.EqualToForwardRef("1 + int", owner=f),
+                "a": int,
+                "b": [
+                    support.EqualToForwardRef("builtins.undef", owner=f),
+                    # We can't resolve this because we have to evaluate the whole annotation
+                    support.EqualToForwardRef("builtins.int", owner=f),
+                ],
+            },
+        )
+
+        self.assertEqual(
+            get_annotations(f, format=Format.STRING),
+            {
+                "x": "builtins.undef",
+                "y": "list[int]",
+                "z": "1 + int",
+                "a": "builtins.int",
+                "b": "[builtins.undef, builtins.int]",
+            },
+        )
+
+    def test_partial_evaluation_cell(self):
+        obj = object()
+
+        class RaisesAttributeError:
+            attriberr: obj.missing
+
+        anno = get_annotations(RaisesAttributeError, format=Format.FORWARDREF)
+        self.assertEqual(
+            anno,
+            {
+                "attriberr": support.EqualToForwardRef(
+                    "obj.missing", is_class=True, owner=RaisesAttributeError
+                )
+            },
+        )
+
 
 class TestCallEvaluateFunction(unittest.TestCase):
     def test_evaluation(self):
@@ -1370,6 +1422,38 @@ class TestForwardRefClass(unittest.TestCase):
             with self.assertRaises(TypeError):
                 pickle.dumps(fr, proto)
 
+    def test_evaluate_string_format(self):
+        fr = ForwardRef("set[Any]")
+        self.assertEqual(fr.evaluate(format=Format.STRING), "set[Any]")
+
+    def test_evaluate_forwardref_format(self):
+        fr = ForwardRef("undef")
+        evaluated = fr.evaluate(format=Format.FORWARDREF)
+        self.assertIs(fr, evaluated)
+
+        fr = ForwardRef("set[undefined]")
+        evaluated = fr.evaluate(format=Format.FORWARDREF)
+        self.assertEqual(
+            evaluated,
+            set[support.EqualToForwardRef("undefined")],
+        )
+
+        fr = ForwardRef("a + b")
+        self.assertEqual(
+            fr.evaluate(format=Format.FORWARDREF),
+            support.EqualToForwardRef("a + b"),
+        )
+        self.assertEqual(
+            fr.evaluate(format=Format.FORWARDREF, locals={"a": 1, "b": 2}),
+            3,
+        )
+
+        fr = ForwardRef('"a" + 1')
+        self.assertEqual(
+            fr.evaluate(format=Format.FORWARDREF),
+            support.EqualToForwardRef('"a" + 1'),
+        )
+
     def test_evaluate_with_type_params(self):
         class Gen[T]:
             alias = int