gh-125614: annotationlib: Fix bug where not all Stringifiers are converted (#125635)

Author: JelleZijlstra

Lib/annotationlib.py

@@ -45,6 +45,7 @@ class Format(enum.IntEnum):
     "__globals__",
     "__owner__",
     "__cell__",
+    "__stringifier_dict__",
 )
 
 
@@ -268,7 +269,16 @@ class _Stringifier:
     # instance of the other in place.
     __slots__ = _SLOTS
 
-    def __init__(self, node, globals=None, owner=None, is_class=False, cell=None):
+    def __init__(
+        self,
+        node,
+        globals=None,
+        owner=None,
+        is_class=False,
+        cell=None,
+        *,
+        stringifier_dict,
+    ):
         # Either an AST node or a simple str (for the common case where a ForwardRef
         # represent a single name).
         assert isinstance(node, (ast.AST, str))
@@ -283,6 +293,7 @@ def __init__(self, node, globals=None, owner=None, is_class=False, cell=None):
         self.__globals__ = globals
         self.__cell__ = cell
         self.__owner__ = owner
+        self.__stringifier_dict__ = stringifier_dict
 
     def __convert_to_ast(self, other):
         if isinstance(other, _Stringifier):
@@ -317,9 +328,15 @@ def __get_ast(self):
         return node
 
     def __make_new(self, node):
-        return _Stringifier(
-            node, self.__globals__, self.__owner__, self.__forward_is_class__
+        stringifier = _Stringifier(
+            node,
+            self.__globals__,
+            self.__owner__,
+            self.__forward_is_class__,
+            stringifier_dict=self.__stringifier_dict__,
         )
+        self.__stringifier_dict__.stringifiers.append(stringifier)
+        return stringifier
 
     # Must implement this since we set __eq__. We hash by identity so that
     # stringifiers in dict keys are kept separate.
@@ -462,6 +479,7 @@ def __missing__(self, key):
             globals=self.globals,
             owner=self.owner,
             is_class=self.is_class,
+            stringifier_dict=self,
         )
         self.stringifiers.append(fwdref)
         return fwdref
@@ -516,7 +534,7 @@ def call_annotate_function(annotate, format, *, owner=None, _is_evaluate=False):
                     name = freevars[i]
                 else:
                     name = "__cell__"
-                fwdref = _Stringifier(name)
+                fwdref = _Stringifier(name, stringifier_dict=globals)
                 new_closure.append(types.CellType(fwdref))
             closure = tuple(new_closure)
         else:
@@ -573,6 +591,7 @@ def call_annotate_function(annotate, format, *, owner=None, _is_evaluate=False):
                         owner=owner,
                         globals=annotate.__globals__,
                         is_class=is_class,
+                        stringifier_dict=globals,
                     )
                     globals.stringifiers.append(fwdref)
                     new_closure.append(types.CellType(fwdref))
@@ -591,6 +610,7 @@ def call_annotate_function(annotate, format, *, owner=None, _is_evaluate=False):
         result = func(Format.VALUE)
         for obj in globals.stringifiers:
             obj.__class__ = ForwardRef
+            obj.__stringifier_dict__ = None  # not needed for ForwardRef
             if isinstance(obj.__ast_node__, str):
                 obj.__arg__ = obj.__ast_node__
                 obj.__ast_node__ = None

Lib/test/test_annotationlib.py

@@ -80,6 +80,42 @@ def f(x: int, y: doesntexist):
             fwdref.evaluate()
         self.assertEqual(fwdref.evaluate(globals={"doesntexist": 1}), 1)
 
+    def test_nonexistent_attribute(self):
+        def f(
+            x: some.module,
+            y: some[module],
+            z: some(module),
+            alpha: some | obj,
+            beta: +some,
+            gamma: some < obj,
+        ):
+            pass
+
+        anno = annotationlib.get_annotations(f, format=Format.FORWARDREF)
+        x_anno = anno["x"]
+        self.assertIsInstance(x_anno, ForwardRef)
+        self.assertEqual(x_anno, ForwardRef("some.module"))
+
+        y_anno = anno["y"]
+        self.assertIsInstance(y_anno, ForwardRef)
+        self.assertEqual(y_anno, ForwardRef("some[module]"))
+
+        z_anno = anno["z"]
+        self.assertIsInstance(z_anno, ForwardRef)
+        self.assertEqual(z_anno, ForwardRef("some(module)"))
+
+        alpha_anno = anno["alpha"]
+        self.assertIsInstance(alpha_anno, ForwardRef)
+        self.assertEqual(alpha_anno, ForwardRef("some | obj"))
+
+        beta_anno = anno["beta"]
+        self.assertIsInstance(beta_anno, ForwardRef)
+        self.assertEqual(beta_anno, ForwardRef("+some"))
+
+        gamma_anno = anno["gamma"]
+        self.assertIsInstance(gamma_anno, ForwardRef)
+        self.assertEqual(gamma_anno, ForwardRef("some < obj"))
+
 
 class TestSourceFormat(unittest.TestCase):
     def test_closure(self):
@@ -91,6 +127,16 @@ def inner(arg: x):
         anno = annotationlib.get_annotations(inner, format=Format.STRING)
         self.assertEqual(anno, {"arg": "x"})
 
+    def test_closure_undefined(self):
+        if False:
+            x = 0
+
+        def inner(arg: x):
+            pass
+
+        anno = annotationlib.get_annotations(inner, format=Format.STRING)
+        self.assertEqual(anno, {"arg": "x"})
+
     def test_function(self):
         def f(x: int, y: doesntexist):
             pass

Misc/NEWS.d/next/Library/2024-10-16-22-45-50.gh-issue-125614.3OEo_Q.rst

@@ -0,0 +1,3 @@
+In the :data:`~annotationlib.Format.FORWARDREF` format of
+:mod:`annotationlib`, fix bug where nested expressions were not returned as
+:class:`annotationlib.ForwardRef` format.