gh-119180: PEP 649: Add __annotate__ attributes (#119209)

Author: JelleZijlstra

Include/cpython/funcobject.h

@@ -41,6 +41,7 @@ typedef struct {
     PyObject *func_weakreflist; /* List of weak references */
     PyObject *func_module;      /* The __module__ attribute, can be anything */
     PyObject *func_annotations; /* Annotations, a dict or NULL */
+    PyObject *func_annotate;    /* Callable to fill the annotations dictionary */
     PyObject *func_typeparams;  /* Tuple of active type variables or NULL */
     vectorcallfunc vectorcall;
     /* Version number for use by specializer.

Include/internal/pycore_global_objects_fini_generated.h

@@ -79,6 +79,7 @@ struct _Py_global_strings {
         STRUCT_FOR_ID(__all__)
         STRUCT_FOR_ID(__and__)
         STRUCT_FOR_ID(__anext__)
+        STRUCT_FOR_ID(__annotate__)
         STRUCT_FOR_ID(__annotations__)
         STRUCT_FOR_ID(__args__)
         STRUCT_FOR_ID(__asyncio_running_event_loop__)

Include/internal/pycore_global_strings.h

@@ -1564,7 +1564,7 @@ def func():
         check(x, size('3Pi2cP7P2ic??2P'))
         # function
         def func(): pass
-        check(func, size('15Pi'))
+        check(func, size('16Pi'))
         class c():
             @staticmethod
             def foo():

Include/internal/pycore_runtime_init_generated.h

@@ -1,4 +1,5 @@
 import textwrap
+import types
 import unittest
 from test.support import run_code
 
@@ -212,3 +213,46 @@ def test_match(self):
                 case 0:
                     x: int = 1
         """)
+
+
+class AnnotateTests(unittest.TestCase):
+    """See PEP 649."""
+    def test_manual_annotate(self):
+        def f():
+            pass
+        mod = types.ModuleType("mod")
+        class X:
+            pass
+
+        for obj in (f, mod, X):
+            with self.subTest(obj=obj):
+                self.check_annotations(obj)
+
+    def check_annotations(self, f):
+        self.assertEqual(f.__annotations__, {})
+        self.assertIs(f.__annotate__, None)
+
+        with self.assertRaisesRegex(TypeError, "__annotate__ must be callable or None"):
+            f.__annotate__ = 42
+        f.__annotate__ = lambda: 42
+        with self.assertRaisesRegex(TypeError, r"takes 0 positional arguments but 1 was given"):
+            print(f.__annotations__)
+
+        f.__annotate__ = lambda x: 42
+        with self.assertRaisesRegex(TypeError, r"__annotate__ returned non-dict of type 'int'"):
+            print(f.__annotations__)
+
+        f.__annotate__ = lambda x: {"x": x}
+        self.assertEqual(f.__annotations__, {"x": 1})
+
+        # Setting annotate to None does not invalidate the cached __annotations__
+        f.__annotate__ = None
+        self.assertEqual(f.__annotations__, {"x": 1})
+
+        # But setting it to a new callable does
+        f.__annotate__ = lambda x: {"y": x}
+        self.assertEqual(f.__annotations__, {"y": 1})
+
+        # Setting f.__annotations__ also clears __annotate__
+        f.__annotations__ = {"z": 43}
+        self.assertIs(f.__annotate__, None)

Include/internal/pycore_unicodeobject_generated.h

@@ -3723,7 +3723,7 @@ def meth(self): pass
 
         acceptable_extra_attrs = {
             '_is_protocol', '_is_runtime_protocol', '__parameters__',
-            '__init__', '__annotations__', '__subclasshook__',
+            '__init__', '__annotations__', '__subclasshook__', '__annotate__',
         }
         self.assertLessEqual(vars(NonP).keys(), vars(C).keys() | acceptable_extra_attrs)
         self.assertLessEqual(

Lib/test/test_sys.py

@@ -1889,6 +1889,7 @@ class _TypingEllipsis:
     '__init__', '__module__', '__new__', '__slots__',
     '__subclasshook__', '__weakref__', '__class_getitem__',
     '__match_args__', '__static_attributes__', '__firstlineno__',
+    '__annotate__',
 })
 
 # These special attributes will be not collected as protocol members.

Lib/test/test_type_annotations.py

@@ -0,0 +1,2 @@
+Add an ``__annotate__`` attribute to functions, classes, and modules as part
+of :pep:`649`. Patch by Jelle Zijlstra.

Lib/test/test_typing.py

@@ -3,6 +3,7 @@
 
 #include "Python.h"
 #include "pycore_ceval.h"         // _PyEval_BuiltinsFromGlobals()
+#include "pycore_long.h"          // _PyLong_GetOne()
 #include "pycore_modsupport.h"    // _PyArg_NoKeywords()
 #include "pycore_object.h"        // _PyObject_GC_UNTRACK()
 #include "pycore_pyerrors.h"      // _PyErr_Occurred()
@@ -124,6 +125,7 @@ _PyFunction_FromConstructor(PyFrameConstructor *constr)
     op->func_weakreflist = NULL;
     op->func_module = module;
     op->func_annotations = NULL;
+    op->func_annotate = NULL;
     op->func_typeparams = NULL;
     op->vectorcall = _PyFunction_Vectorcall;
     op->func_version = 0;
@@ -202,6 +204,7 @@ PyFunction_NewWithQualName(PyObject *code, PyObject *globals, PyObject *qualname
     op->func_weakreflist = NULL;
     op->func_module = module;
     op->func_annotations = NULL;
+    op->func_annotate = NULL;
     op->func_typeparams = NULL;
     op->vectorcall = _PyFunction_Vectorcall;
     op->func_version = 0;
@@ -512,7 +515,22 @@ static PyObject *
 func_get_annotation_dict(PyFunctionObject *op)
 {
     if (op->func_annotations == NULL) {
-        return NULL;
+        if (op->func_annotate == NULL || !PyCallable_Check(op->func_annotate)) {
+            Py_RETURN_NONE;
+        }
+        PyObject *one = _PyLong_GetOne();
+        PyObject *ann_dict = _PyObject_CallOneArg(op->func_annotate, one);
+        if (ann_dict == NULL) {
+            return NULL;
+        }
+        if (!PyDict_Check(ann_dict)) {
+            PyErr_Format(PyExc_TypeError, "__annotate__ returned non-dict of type '%.100s'",
+                         Py_TYPE(ann_dict)->tp_name);
+            Py_DECREF(ann_dict);
+            return NULL;
+        }
+        Py_XSETREF(op->func_annotations, ann_dict);
+        return ann_dict;
     }
     if (PyTuple_CheckExact(op->func_annotations)) {
         PyObject *ann_tuple = op->func_annotations;
@@ -565,7 +583,9 @@ PyFunction_SetAnnotations(PyObject *op, PyObject *annotations)
                         "non-dict annotations");
         return -1;
     }
-    Py_XSETREF(((PyFunctionObject *)op)->func_annotations, annotations);
+    PyFunctionObject *func = (PyFunctionObject *)op;
+    Py_XSETREF(func->func_annotations, annotations);
+    Py_CLEAR(func->func_annotate);
     return 0;
 }
 
@@ -763,10 +783,44 @@ func_set_kwdefaults(PyFunctionObject *op, PyObject *value, void *Py_UNUSED(ignor
     return 0;
 }
 
+static PyObject *
+func_get_annotate(PyFunctionObject *op, void *Py_UNUSED(ignored))
+{
+    if (op->func_annotate == NULL) {
+        Py_RETURN_NONE;
+    }
+    return Py_NewRef(op->func_annotate);
+}
+
+static int
+func_set_annotate(PyFunctionObject *op, PyObject *value, void *Py_UNUSED(ignored))
+{
+    if (value == NULL) {
+        PyErr_SetString(PyExc_TypeError,
+            "__annotate__ cannot be deleted");
+        return -1;
+    }
+    if (Py_IsNone(value)) {
+        Py_XSETREF(op->func_annotate, value);
+        return 0;
+    }
+    else if (PyCallable_Check(value)) {
+        Py_XSETREF(op->func_annotate, Py_XNewRef(value));
+        Py_CLEAR(op->func_annotations);
+        return 0;
+    }
+    else {
+        PyErr_SetString(PyExc_TypeError,
+            "__annotate__ must be callable or None");
+        return -1;
+    }
+}
+
 static PyObject *
 func_get_annotations(PyFunctionObject *op, void *Py_UNUSED(ignored))
 {
-    if (op->func_annotations == NULL) {
+    if (op->func_annotations == NULL &&
+        (op->func_annotate == NULL || !PyCallable_Check(op->func_annotate))) {
         op->func_annotations = PyDict_New();
         if (op->func_annotations == NULL)
             return NULL;
@@ -789,6 +843,7 @@ func_set_annotations(PyFunctionObject *op, PyObject *value, void *Py_UNUSED(igno
         return -1;
     }
     Py_XSETREF(op->func_annotations, Py_XNewRef(value));
+    Py_CLEAR(op->func_annotate);
     return 0;
 }
 
@@ -836,6 +891,7 @@ static PyGetSetDef func_getsetlist[] = {
      (setter)func_set_kwdefaults},
     {"__annotations__", (getter)func_get_annotations,
      (setter)func_set_annotations},
+    {"__annotate__", (getter)func_get_annotate, (setter)func_set_annotate},
     {"__dict__", PyObject_GenericGetDict, PyObject_GenericSetDict},
     {"__name__", (getter)func_get_name, (setter)func_set_name},
     {"__qualname__", (getter)func_get_qualname, (setter)func_set_qualname},
@@ -972,6 +1028,7 @@ func_clear(PyFunctionObject *op)
     Py_CLEAR(op->func_dict);
     Py_CLEAR(op->func_closure);
     Py_CLEAR(op->func_annotations);
+    Py_CLEAR(op->func_annotate);
     Py_CLEAR(op->func_typeparams);
     // Don't Py_CLEAR(op->func_code), since code is always required
     // to be non-NULL. Similarly, name and qualname shouldn't be NULL.
@@ -1028,6 +1085,7 @@ func_traverse(PyFunctionObject *f, visitproc visit, void *arg)
     Py_VISIT(f->func_dict);
     Py_VISIT(f->func_closure);
     Py_VISIT(f->func_annotations);
+    Py_VISIT(f->func_annotate);
     Py_VISIT(f->func_typeparams);
     Py_VISIT(f->func_qualname);
     return 0;