diff --git a/Include/cpython/funcobject.h b/Include/cpython/funcobject.h index 598cd330bc9ca9..2c9be13ab658b0 100644 --- a/Include/cpython/funcobject.h +++ b/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. diff --git a/Lib/test/test_dis.py b/Lib/test/test_dis.py index da57aad2f84fbd..a2d129db9dbc31 100644 --- a/Lib/test/test_dis.py +++ b/Lib/test/test_dis.py @@ -411,7 +411,6 @@ def foo(a: int, b: str) -> str: 0 RESUME 0 2 LOAD_CONST 0 (", line 2>) - MAKE_FUNCTION LOAD_CONST 1 (", line 2>) MAKE_FUNCTION SET_FUNCTION_ATTRIBUTE 16 (annotate) diff --git a/Objects/funcobject.c b/Objects/funcobject.c index 7b17a9ba31fac4..d71e604c4c1bb2 100644 --- a/Objects/funcobject.c +++ b/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,49 @@ 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)) { + PyObject *func = PyFunction_New(op->func_annotate, op->func_globals); + if (func == NULL) { + return NULL; + } + Py_SETREF(op->func_annotate, Py_NewRef(func)); + return func; + } + 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; + Py_SETREF(op->func_annotate, Py_NewRef(func)); + 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 +907,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; diff --git a/Python/codegen.c b/Python/codegen.c index 61707ba677097c..77f460e23f26f0 100644 --- a/Python/codegen.c +++ b/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 */ @@ -690,7 +692,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); @@ -726,7 +729,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; @@ -774,7 +778,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; @@ -846,6 +851,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) @@ -1059,7 +1090,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; }