Add tests for replacement of extra names in repr and string evaluation

Author: DavidCEllis

Lib/test/test_annotationlib.py

@@ -1961,6 +1961,11 @@ def test_forward_repr(self):
             "typing.List[ForwardRef('int', owner='class')]",
         )
 
+    def test_forward_repr_extra_names(self):
+        fr = ForwardRef("__annotationlib_name_1__")
+        fr.__extra_names__ = {"__annotationlib_name_1__": list[str]}
+        self.assertEqual(repr(fr), "ForwardRef('list[str]')")
+
     def test_forward_recursion_actually(self):
         def namespace1():
             a = ForwardRef("A")
@@ -2037,6 +2042,20 @@ def test_evaluate_string_format(self):
         fr = ForwardRef("set[Any]")
         self.assertEqual(fr.evaluate(format=Format.STRING), "set[Any]")
 
+    def test_evaluate_string_format_extra_names(self):
+        # Test that internal extra_names are replaced when evaluating as strings
+
+        # As identifier
+        fr = ForwardRef("__annotationlib_name_1__")
+        fr.__extra_names__ = {"__annotationlib_name_1__": str}
+        self.assertEqual(fr.evaluate(format=Format.STRING), "str")
+
+        # Via AST visitor
+        def f(a: ref | str): ...
+
+        fr = get_annotations(f, format=Format.FORWARDREF)['a']
+        self.assertEqual(fr.evaluate(format=Format.STRING), "ref | str")
+
     def test_evaluate_forwardref_format(self):
         fr = ForwardRef("undef")
         evaluated = fr.evaluate(format=Format.FORWARDREF)