Use the same annotation logic for VALUE, FORWARDREF and STRING

Also split out the single large test into multiple smaller tests

Author: DavidCEllis

Lib/dataclasses.py

@@ -537,36 +537,34 @@ def _make_annotate_function(cls, annotations):
     # annotations should be in FORWARDREF format at this stage
 
     def __annotate__(format, /):
+        Format = annotationlib.Format
         match format:
-            case annotationlib.Format.VALUE | annotationlib.Format.FORWARDREF:
-                return {
-                    k: v.evaluate(format=format)
-                    if isinstance(v, annotationlib.ForwardRef) else v
-                    for k, v in annotations.items()
-                }
-
-            case annotationlib.Format.STRING:
+            case Format.VALUE | Format.FORWARDREF | Format.STRING:
                 cls_annotations = {}
                 for base in reversed(cls.__mro__):
                     cls_annotations.update(
                         annotationlib.get_annotations(base, format=format)
                     )
 
-                string_annos = {}
+                new_annotations = {}
                 for k, v in annotations.items():
                     try:
-                        string_annos[k] = cls_annotations[k]
+                        new_annotations[k] = cls_annotations[k]
                     except KeyError:
                         # This should be the return value
-                        string_annos[k] = annotationlib.type_repr(v)
-                return string_annos
+                        if format == Format.STRING:
+                            new_annotations[k] = annotationlib.type_repr(v)
+                        else:
+                            new_annotations[k] = v
+
+                return new_annotations
 
             case _:
                 raise NotImplementedError(format)
 
     # This is a flag for _add_slots to know it needs to regenerate this method
     # In order to remove references to the original class when it is replaced
-    __annotate__.__generated_by_dataclasses = True
+    __annotate__._generated_by_dataclasses = True
 
     return __annotate__
 
@@ -1399,7 +1397,7 @@ def _add_slots(cls, is_frozen, weakref_slot, defined_fields):
 
     # Fix references in generated __annotate__ methods
     method = getattr(newcls, "__init__")
-    update_annotations = getattr(method.__annotate__, "__generated_by_dataclasses", False)
+    update_annotations = getattr(method.__annotate__, "_generated_by_dataclasses", False)
 
     if update_annotations:
         new_annotations = {}

Lib/test/test_dataclasses/__init__.py

@@ -2470,10 +2470,13 @@ def __init__(self, a):
 
         self.assertEqual(D(5).a, 10)
 
+
+class TestInitAnnotate(unittest.TestCase):
+    # Tests for the generated __annotate__ function for __init__
+    # See: https://github.com/python/cpython/issues/137530
+
     def test_annotate_function(self):
-        # Test that the __init__ function has correct annotate function
-        # See: https://github.com/python/cpython/issues/137530
-        # With no forward references
+        # No forward references
         @dataclass
         class A:
             a: int
@@ -2486,6 +2489,9 @@ class A:
         self.assertEqual(forwardref_annos, {'a': int, 'return': None})
         self.assertEqual(string_annos, {'a': 'int', 'return': 'None'})
 
+        self.assertTrue(getattr(A.__init__.__annotate__, "_generated_by_dataclasses"))
+
+    def test_annotate_function_forwardref(self):
         # With forward references
         @dataclass
         class B:
@@ -2512,26 +2518,53 @@ class B:
         self.assertEqual(forwardref_annos, {'b': int, 'return': None})
         self.assertEqual(string_annos, {'b': 'undefined', 'return': 'None'})
 
-        del undefined  # Remove so we can use the name in later examples
-
+    def test_annotate_function_init_false(self):
         # Check `init=False` attributes don't get into the annotations of the __init__ function
         @dataclass
         class C:
             c: str = field(init=False)
 
         self.assertEqual(annotationlib.get_annotations(C.__init__), {'return': None})
 
-
+    def test_annotate_function_contains_forwardref(self):
         # Check string annotations on objects containing a ForwardRef
         @dataclass
         class D:
             d: list[undefined]
 
+        with self.assertRaises(NameError):
+            annotationlib.get_annotations(D.__init__)
+
+        self.assertEqual(
+            annotationlib.get_annotations(D.__init__, format=annotationlib.Format.FORWARDREF),
+            {"d": list[support.EqualToForwardRef("undefined", is_class=True, owner=D)], "return": None}
+        )
+
+        self.assertEqual(
+            annotationlib.get_annotations(D.__init__, format=annotationlib.Format.STRING),
+            {"d": "list[undefined]", "return": "None"}
+        )
+
+        # Now test when it is defined
+        undefined = str
+
+        # VALUE should now resolve
+        self.assertEqual(
+            annotationlib.get_annotations(D.__init__),
+            {"d": list[str], "return": None}
+        )
+
+        self.assertEqual(
+            annotationlib.get_annotations(D.__init__, format=annotationlib.Format.FORWARDREF),
+            {"d": list[str], "return": None}
+        )
+
         self.assertEqual(
             annotationlib.get_annotations(D.__init__, format=annotationlib.Format.STRING),
             {"d": "list[undefined]", "return": "None"}
         )
 
+    def test_annotate_function_not_replaced(self):
         # Check that __annotate__ is not replaced on non-generated __init__ functions
         @dataclass(slots=True)
         class E:
@@ -2543,6 +2576,26 @@ def __init__(self, x: int) -> None:
             annotationlib.get_annotations(E.__init__), {"x": int, "return": None}
         )
 
+        self.assertFalse(hasattr(E.__init__.__annotate__, "_generated_by_dataclasses"))
+
+    def test_init_false_forwardref(self):
+        # Currently this raises a NameError even though the ForwardRef
+        # is not in the __init__ method
+
+        @dataclass
+        class F:
+            not_in_init: list[undefined] = field(init=False, default=None)
+            in_init: int
+
+        annos = annotationlib.get_annotations(F.__init__, format=annotationlib.Format.FORWARDREF)
+        self.assertEqual(
+            annos,
+            {"in_init": int, "return": None},
+        )
+
+        with self.assertRaises(NameError):
+            annos = annotationlib.get_annotations(F.__init__)  # NameError on not_in_init
+
 
 class TestRepr(unittest.TestCase):
     def test_repr(self):