fix tests

JelleZijlstra · Jan 17, 2025 · a4ed192 · a4ed192
1 parent a245727
commit a4ed192
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 2 deletions.
diff --git 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 (<code object __annotate__ at 0x..., file "<dis>", line 2>)
-              MAKE_FUNCTION
               LOAD_CONST               1 (<code object foo at 0x..., file "<dis>", line 2>)
               MAKE_FUNCTION
               SET_FUNCTION_ATTRIBUTE  16 (annotate)

diff --git a/Objects/funcobject.c b/Objects/funcobject.c
@@ -845,7 +845,12 @@ get_annotate_function(PyFunctionObject *op)
         return Py_NewRef(op->func_annotate);
     }
     else if (PyCode_Check(op->func_annotate)) {
-        return PyFunction_New(op->func_annotate, op->func_globals);
+        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);
@@ -863,6 +868,7 @@ get_annotate_function(PyFunctionObject *op)
             return NULL;
         }
         _PyFunction_CAST(func)->func_closure = closure;
+        Py_SETREF(op->func_annotate, Py_NewRef(func));
         return func;
     }
     PyErr_Format(PyExc_SystemError,