gh-121332: Make AST node constructor check _attributes instead of hardcoding attributes (#121334)

JelleZijlstra · web-flow · commit 58e8cf2bb61f · 2024-07-11T14:34:53.000Z
diff --git a/Lib/test/test_ast.py b/Lib/test/test_ast.py
@@ -1386,15 +1386,7 @@ class MyNode(ast.AST):
         self.assertEqual(node.y, 1)
 
         y = object()
-        # custom attributes are currently not supported and raise a warning
-        # because the allowed attributes are hard-coded !
-        msg = (
-            "MyNode.__init__ got an unexpected keyword argument 'y'. "
-            "Support for arbitrary keyword arguments is deprecated and "
-            "will be removed in Python 3.15"
-        )
-        with self.assertWarnsRegex(DeprecationWarning, re.escape(msg)):
-            repl = copy.replace(node, y=y)
+        repl = copy.replace(node, y=y)
         # assert that there is no side-effect
         self.assertEqual(node.x, 0)
         self.assertEqual(node.y, 1)
@@ -3250,6 +3242,18 @@ class FieldsAndTypes(ast.AST):
         obj = FieldsAndTypes(a=1)
         self.assertEqual(obj.a, 1)
 
+    def test_custom_attributes(self):
+        class MyAttrs(ast.AST):
+            _attributes = ("a", "b")
+
+        obj = MyAttrs(a=1, b=2)
+        self.assertEqual(obj.a, 1)
+        self.assertEqual(obj.b, 2)
+
+        with self.assertWarnsRegex(DeprecationWarning,
+                                   r"MyAttrs.__init__ got an unexpected keyword argument 'c'."):
+            obj = MyAttrs(c=3)
+
     def test_fields_and_types_no_default(self):
         class FieldsAndTypesNoDefault(ast.AST):
             _fields = ('a',)
diff --git a/Misc/NEWS.d/next/Library/2024-07-03-07-25-21.gh-issue-121332.Iz6FEq.rst b/Misc/NEWS.d/next/Library/2024-07-03-07-25-21.gh-issue-121332.Iz6FEq.rst
@@ -0,0 +1,4 @@
+Fix constructor of :mod:`ast` nodes with custom ``_attributes``. Previously,
+passing custom attributes would raise a :py:exc:`DeprecationWarning`. Passing
+arguments to the constructor that are not in ``_fields`` or ``_attributes``
+remains deprecated. Patch by Jelle Zijlstra.
diff --git a/Parser/asdl_c.py b/Parser/asdl_c.py
@@ -880,7 +880,7 @@ def visitModule(self, mod):
 
     Py_ssize_t i, numfields = 0;
     int res = -1;
-    PyObject *key, *value, *fields, *remaining_fields = NULL;
+    PyObject *key, *value, *fields, *attributes = NULL, *remaining_fields = NULL;
     if (PyObject_GetOptionalAttr((PyObject*)Py_TYPE(self), state->_fields, &fields) < 0) {
         goto cleanup;
     }
@@ -947,22 +947,32 @@ def visitModule(self, mod):
                     goto cleanup;
                 }
             }
-            else if (
-                PyUnicode_CompareWithASCIIString(key, "lineno") != 0 &&
-                PyUnicode_CompareWithASCIIString(key, "col_offset") != 0 &&
-                PyUnicode_CompareWithASCIIString(key, "end_lineno") != 0 &&
-                PyUnicode_CompareWithASCIIString(key, "end_col_offset") != 0
-            ) {
-                if (PyErr_WarnFormat(
-                    PyExc_DeprecationWarning, 1,
-                    "%.400s.__init__ got an unexpected keyword argument '%U'. "
-                    "Support for arbitrary keyword arguments is deprecated "
-                    "and will be removed in Python 3.15.",
-                    Py_TYPE(self)->tp_name, key
-                ) < 0) {
+            else {
+                // Lazily initialize "attributes"
+                if (attributes == NULL) {
+                    attributes = PyObject_GetAttr((PyObject*)Py_TYPE(self), state->_attributes);
+                    if (attributes == NULL) {
+                        res = -1;
+                        goto cleanup;
+                    }
+                }
+                int contains = PySequence_Contains(attributes, key);
+                if (contains == -1) {
                     res = -1;
                     goto cleanup;
                 }
+                else if (contains == 0) {
+                    if (PyErr_WarnFormat(
+                        PyExc_DeprecationWarning, 1,
+                        "%.400s.__init__ got an unexpected keyword argument '%U'. "
+                        "Support for arbitrary keyword arguments is deprecated "
+                        "and will be removed in Python 3.15.",
+                        Py_TYPE(self)->tp_name, key
+                    ) < 0) {
+                        res = -1;
+                        goto cleanup;
+                    }
+                }
             }
             res = PyObject_SetAttr(self, key, value);
             if (res < 0) {
@@ -1045,6 +1055,7 @@ def visitModule(self, mod):
         Py_DECREF(field_types);
     }
   cleanup:
+    Py_XDECREF(attributes);
     Py_XDECREF(fields);
     Py_XDECREF(remaining_fields);
     return res;
diff --git a/Python/Python-ast.c b/Python/Python-ast.c
