gh-119180: Lazily wrap annotations on classmethod and staticmethod (#119864)

JelleZijlstra · web-flow · commit d28afd3fa064 · 2024-05-31T14:05:51.000-07:00
diff --git a/Lib/test/test_descr.py b/Lib/test/test_descr.py
@@ -1593,8 +1593,7 @@ def f(cls, arg):
             self.fail("classmethod shouldn't accept keyword args")
 
         cm = classmethod(f)
-        cm_dict = {'__annotations__': {},
-                   '__doc__': (
+        cm_dict = {'__doc__': (
                        "f docstring"
                        if support.HAVE_DOCSTRINGS
                        else None
@@ -1610,6 +1609,41 @@ def f(cls, arg):
         del cm.x
         self.assertNotHasAttr(cm, "x")
 
+    def test_classmethod_staticmethod_annotations(self):
+        for deco in (classmethod, staticmethod):
+            @deco
+            def unannotated(cls): pass
+            @deco
+            def annotated(cls) -> int: pass
+
+            for method in (annotated, unannotated):
+                with self.subTest(deco=deco, method=method):
+                    original_annotations = dict(method.__wrapped__.__annotations__)
+                    self.assertNotIn('__annotations__', method.__dict__)
+                    self.assertEqual(method.__annotations__, original_annotations)
+                    self.assertIn('__annotations__', method.__dict__)
+
+                    new_annotations = {"a": "b"}
+                    method.__annotations__ = new_annotations
+                    self.assertEqual(method.__annotations__, new_annotations)
+                    self.assertEqual(method.__wrapped__.__annotations__, original_annotations)
+
+                    del method.__annotations__
+                    self.assertEqual(method.__annotations__, original_annotations)
+
+                    original_annotate = method.__wrapped__.__annotate__
+                    self.assertNotIn('__annotate__', method.__dict__)
+                    self.assertIs(method.__annotate__, original_annotate)
+                    self.assertIn('__annotate__', method.__dict__)
+
+                    new_annotate = lambda: {"annotations": 1}
+                    method.__annotate__ = new_annotate
+                    self.assertIs(method.__annotate__, new_annotate)
+                    self.assertIs(method.__wrapped__.__annotate__, original_annotate)
+
+                    del method.__annotate__
+                    self.assertIs(method.__annotate__, original_annotate)
+
     @support.refcount_test
     def test_refleaks_in_classmethod___init__(self):
         gettotalrefcount = support.get_attribute(sys, 'gettotalrefcount')
diff --git a/Misc/NEWS.d/next/Core and Builtins/2024-05-31-08-23-41.gh-issue-119180.KL4VxZ.rst b/Misc/NEWS.d/next/Core and Builtins/2024-05-31-08-23-41.gh-issue-119180.KL4VxZ.rst
@@ -0,0 +1,3 @@
+:func:`classmethod` and :func:`staticmethod` now wrap the
+:attr:`__annotations__` and :attr:`!__annotate__` attributes of their
+underlying callable lazily. See :pep:`649`. Patch by Jelle Zijlstra.
diff --git a/Objects/funcobject.c b/Objects/funcobject.c
@@ -1172,12 +1172,57 @@ functools_wraps(PyObject *wrapper, PyObject *wrapped)
     COPY_ATTR(__name__);
     COPY_ATTR(__qualname__);
     COPY_ATTR(__doc__);
-    COPY_ATTR(__annotations__);
     return 0;
 
 #undef COPY_ATTR
 }
 
+// Used for wrapping __annotations__ and __annotate__ on classmethod
+// and staticmethod objects.
+static PyObject *
+descriptor_get_wrapped_attribute(PyObject *wrapped, PyObject *dict, PyObject *name)
+{
+    PyObject *res;
+    if (PyDict_GetItemRef(dict, name, &res) < 0) {
+        return NULL;
+    }
+    if (res != NULL) {
+        return res;
+    }
+    res = PyObject_GetAttr(wrapped, name);
+    if (res == NULL) {
+        return NULL;
+    }
+    if (PyDict_SetItem(dict, name, res) < 0) {
+        Py_DECREF(res);
+        return NULL;
+    }
+    return res;
+}
+
+static int
+descriptor_set_wrapped_attribute(PyObject *dict, PyObject *name, PyObject *value,
+                                 char *type_name)
+{
+    if (value == NULL) {
+        if (PyDict_DelItem(dict, name) < 0) {
+            if (PyErr_ExceptionMatches(PyExc_KeyError)) {
+                PyErr_Clear();
+                PyErr_Format(PyExc_AttributeError,
+                             "'%.200s' object has no attribute '%U'",
+                             type_name, name);
+            }
+            else {
+                return -1;
+            }
+        }
+        return 0;
+    }
+    else {
+        return PyDict_SetItem(dict, name, value);
+    }
+}
+
 
 /* Class method object */
 
@@ -1283,10 +1328,37 @@ cm_get___isabstractmethod__(classmethod *cm, void *closure)
     Py_RETURN_FALSE;
 }
 
+static PyObject *
+cm_get___annotations__(classmethod *cm, void *closure)
+{
+    return descriptor_get_wrapped_attribute(cm->cm_callable, cm->cm_dict, &_Py_ID(__annotations__));
+}
+
+static int
+cm_set___annotations__(classmethod *cm, PyObject *value, void *closure)
+{
+    return descriptor_set_wrapped_attribute(cm->cm_dict, &_Py_ID(__annotations__), value, "classmethod");
+}
+
+static PyObject *
+cm_get___annotate__(classmethod *cm, void *closure)
+{
+    return descriptor_get_wrapped_attribute(cm->cm_callable, cm->cm_dict, &_Py_ID(__annotate__));
+}
+
+static int
+cm_set___annotate__(classmethod *cm, PyObject *value, void *closure)
+{
+    return descriptor_set_wrapped_attribute(cm->cm_dict, &_Py_ID(__annotate__), value, "classmethod");
+}
+
+
 static PyGetSetDef cm_getsetlist[] = {
     {"__isabstractmethod__",
      (getter)cm_get___isabstractmethod__, NULL, NULL, NULL},
     {"__dict__", PyObject_GenericGetDict, PyObject_GenericSetDict, NULL, NULL},
+    {"__annotations__", (getter)cm_get___annotations__, (setter)cm_set___annotations__, NULL, NULL},
+    {"__annotate__", (getter)cm_get___annotate__, (setter)cm_set___annotate__, NULL, NULL},
     {NULL} /* Sentinel */
 };
 
@@ -1479,10 +1551,36 @@ sm_get___isabstractmethod__(staticmethod *sm, void *closure)
     Py_RETURN_FALSE;
 }
 
+static PyObject *
+sm_get___annotations__(staticmethod *sm, void *closure)
+{
+    return descriptor_get_wrapped_attribute(sm->sm_callable, sm->sm_dict, &_Py_ID(__annotations__));
+}
+
+static int
+sm_set___annotations__(staticmethod *sm, PyObject *value, void *closure)
+{
+    return descriptor_set_wrapped_attribute(sm->sm_dict, &_Py_ID(__annotations__), value, "staticmethod");
+}
+
+static PyObject *
+sm_get___annotate__(staticmethod *sm, void *closure)
+{
+    return descriptor_get_wrapped_attribute(sm->sm_callable, sm->sm_dict, &_Py_ID(__annotate__));
+}
+
+static int
+sm_set___annotate__(staticmethod *sm, PyObject *value, void *closure)
+{
+    return descriptor_set_wrapped_attribute(sm->sm_dict, &_Py_ID(__annotate__), value, "staticmethod");
+}
+
 static PyGetSetDef sm_getsetlist[] = {
     {"__isabstractmethod__",
      (getter)sm_get___isabstractmethod__, NULL, NULL, NULL},
     {"__dict__", PyObject_GenericGetDict, PyObject_GenericSetDict, NULL, NULL},
+    {"__annotations__", (getter)sm_get___annotations__, (setter)sm_set___annotations__, NULL, NULL},
+    {"__annotate__", (getter)sm_get___annotate__, (setter)sm_set___annotate__, NULL, NULL},
     {NULL} /* Sentinel */
 };
 
