gh-124157: Lazily create __annotate__ functions for function annotations

Author: JelleZijlstra

Include/cpython/funcobject.h

@@ -41,7 +41,14 @@ 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 */
+    /* Callable to fill the annotations dictionary.
+     * May also be:
+     *      NULL (function has no annotations, or only an annotations dict)
+     *      None (set manually by user)
+     *      a code object (if the function has no closure)
+     *      a tuple containing a code object plus cell variables (for a closure)
+     */
+    PyObject *func_annotate;
     PyObject *func_typeparams;  /* Tuple of active type variables or NULL */
     vectorcallfunc vectorcall;
     /* Version number for use by specializer.

Objects/funcobject.c

@@ -8,6 +8,8 @@
 #include "pycore_object.h"        // _PyObject_GC_UNTRACK()
 #include "pycore_pyerrors.h"      // _PyErr_Occurred()
 
+static PyObject *
+get_annotate_function(PyFunctionObject *func);
 
 static const char *
 func_event_name(PyFunction_WatchEvent event) {
@@ -547,11 +549,13 @@ static PyObject *
 func_get_annotation_dict(PyFunctionObject *op)
 {
     if (op->func_annotations == NULL) {
-        if (op->func_annotate == NULL || !PyCallable_Check(op->func_annotate)) {
+        if (op->func_annotate == NULL || Py_IsNone(op->func_annotate)) {
             Py_RETURN_NONE;
         }
+        PyObject *annotate = get_annotate_function(op);
         PyObject *one = _PyLong_GetOne();
-        PyObject *ann_dict = _PyObject_CallOneArg(op->func_annotate, one);
+        PyObject *ann_dict = _PyObject_CallOneArg(annotate, one);
+        Py_DECREF(annotate);
         if (ann_dict == NULL) {
             return NULL;
         }
@@ -828,10 +832,43 @@ static PyObject *
 func_get_annotate(PyObject *self, void *Py_UNUSED(ignored))
 {
     PyFunctionObject *op = _PyFunction_CAST(self);
-    if (op->func_annotate == NULL) {
+    return get_annotate_function(op);
+}
+
+static PyObject *
+get_annotate_function(PyFunctionObject *op)
+{
+    if (op->func_annotate == NULL || Py_IsNone(op->func_annotate)) {
         Py_RETURN_NONE;
     }
-    return Py_NewRef(op->func_annotate);
+    if (PyCallable_Check(op->func_annotate)) {
+        return Py_NewRef(op->func_annotate);
+    }
+    else if (PyCode_Check(op->func_annotate)) {
+        return PyFunction_New(op->func_annotate, op->func_globals);
+    }
+    else if (PyTuple_CheckExact(op->func_annotate) && PyTuple_GET_SIZE(op->func_annotate) >= 2) {
+        PyObject *co = PyTuple_GET_ITEM(op->func_annotate, 0);
+        if (!PyCode_Check(co)) {
+            PyErr_Format(PyExc_SystemError,
+                "func_annotate tuple should contain code object, not '%.100s'",
+                Py_TYPE(co)->tp_name);
+            return NULL;
+        }
+        PyObject *func = PyFunction_New(co, op->func_globals);
+        PyObject *closure = PyTuple_GetSlice(
+            op->func_annotate, 1, PyTuple_GET_SIZE(op->func_annotate));
+        if (closure == NULL) {
+            Py_DECREF(func);
+            return NULL;
+        }
+        _PyFunction_CAST(func)->func_closure = closure;
+        return func;
+    }
+    PyErr_Format(PyExc_SystemError,
+        "Invalid func_annotate attribute of type '%.100s'",
+        Py_TYPE(op->func_annotate)->tp_name);
+    return NULL;
 }
 
 static int
@@ -864,7 +901,7 @@ func_get_annotations(PyObject *self, void *Py_UNUSED(ignored))
 {
     PyFunctionObject *op = _PyFunction_CAST(self);
     if (op->func_annotations == NULL &&
-        (op->func_annotate == NULL || !PyCallable_Check(op->func_annotate))) {
+        (op->func_annotate == NULL || Py_IsNone(op->func_annotate))) {
         op->func_annotations = PyDict_New();
         if (op->func_annotations == NULL)
             return NULL;

Python/codegen.c

@@ -242,6 +242,8 @@ static int codegen_pattern_subpattern(compiler *,
                                       pattern_ty, pattern_context *);
 static int codegen_make_closure(compiler *c, location loc,
                                 PyCodeObject *co, Py_ssize_t flags);
+static int codegen_make_annotate_function(compiler *c, location loc,
+                                          PyCodeObject *co);
 
 
 /* Add an opcode with an integer argument */
@@ -696,7 +698,8 @@ codegen_setup_annotations_scope(compiler *c, location loc,
 
 static int
 codegen_leave_annotations_scope(compiler *c, location loc,
-                                Py_ssize_t annotations_len)
+                                Py_ssize_t annotations_len,
+                                bool is_function)
 {
     ADDOP_I(c, loc, BUILD_MAP, annotations_len);
     ADDOP_IN_SCOPE(c, loc, RETURN_VALUE);
@@ -732,7 +735,8 @@ codegen_leave_annotations_scope(compiler *c, location loc,
     if (co == NULL) {
         return ERROR;
     }
-    int ret = codegen_make_closure(c, loc, co, 0);
+    int ret = is_function ? codegen_make_annotate_function(c, loc, co)
+                          : codegen_make_closure(c, loc, co, 0);
     Py_DECREF(co);
     RETURN_IF_ERROR(ret);
     return SUCCESS;
@@ -780,7 +784,8 @@ codegen_process_deferred_annotations(compiler *c, location loc)
     }
     Py_DECREF(deferred_anno);
 
-    RETURN_IF_ERROR(codegen_leave_annotations_scope(c, loc, annotations_len));
+    RETURN_IF_ERROR(codegen_leave_annotations_scope(c, loc, annotations_len,
+                    /* is_function */false));
     RETURN_IF_ERROR(codegen_nameop(c, loc, &_Py_ID(__annotate__), Store));
 
     return SUCCESS;
@@ -852,6 +857,32 @@ _PyCodegen_EnterAnonymousScope(compiler* c, mod_ty mod)
     return SUCCESS;
 }
 
+static int
+codegen_make_annotate_function(compiler *c, location loc,
+                               PyCodeObject *co)
+{
+    // For annotate functions, we don't bother creating the function object
+    // unless the __annotate__ attribute is accessed.
+    // We may store the prepared data as follows:
+    // - Just a code object
+    // - A tuple containing the code object, followed by a number of cell objects.
+    ADDOP_LOAD_CONST(c, loc, (PyObject*)co);
+    if (co->co_nfreevars) {
+        int i = PyUnstable_Code_GetFirstFree(co);
+        for (; i < co->co_nlocalsplus; ++i) {
+            /* Bypass com_addop_varname because it will generate
+               LOAD_DEREF but LOAD_CLOSURE is needed.
+            */
+            PyObject *name = PyTuple_GET_ITEM(co->co_localsplusnames, i);
+            int arg = _PyCompile_LookupArg(c, co, name);
+            RETURN_IF_ERROR(arg);
+            ADDOP_I(c, loc, LOAD_CLOSURE, arg);
+        }
+        ADDOP_I(c, loc, BUILD_TUPLE, co->co_nfreevars + 1);
+    }
+    return SUCCESS;
+}
+
 static int
 codegen_make_closure(compiler *c, location loc,
                      PyCodeObject *co, Py_ssize_t flags)
@@ -1065,7 +1096,7 @@ codegen_annotations(compiler *c, location loc,
             c, codegen_annotations_in_scope(c, loc, args, returns, &annotations_len)
         );
         RETURN_IF_ERROR(
-            codegen_leave_annotations_scope(c, loc, annotations_len)
+            codegen_leave_annotations_scope(c, loc, annotations_len, /* is_function */true)
         );
         return MAKE_FUNCTION_ANNOTATE;
     }