pythongh-118473: Fix set_asyncgen_hooks not to allow partial setup

youknowone · May 1, 2024 · c8474fc · c8474fc
1 parent c1bf487
commit c8474fc
Show file tree

Hide file tree

Showing 5 changed files with 45 additions and 22 deletions.
diff --git a/Include/internal/pycore_ceval.h b/Include/internal/pycore_ceval.h
@@ -71,8 +71,8 @@ extern PyObject* _PyEval_GetAsyncGenFirstiter(void);
 extern PyObject* _PyEval_GetAsyncGenFinalizer(void);
 
 // Used by sys.set_asyncgen_hooks()
-extern int _PyEval_SetAsyncGenFirstiter(PyObject *);
-extern int _PyEval_SetAsyncGenFinalizer(PyObject *);
+extern void _PyEval_SetAsyncGenFirstiter(PyObject *);
+extern void _PyEval_SetAsyncGenFinalizer(PyObject *);
 
 // Used by sys.get_coroutine_origin_tracking_depth()
 // and sys.set_coroutine_origin_tracking_depth()

diff --git a/Lib/test/test_sys.py b/Lib/test/test_sys.py
@@ -1779,14 +1779,34 @@ def test_asyncgen_hooks(self):
         self.assertIsNone(old.finalizer)
 
         firstiter = lambda *a: None
+        finalizer = lambda *a: None
+
+        try:
+            sys.set_asyncgen_hooks(firstiter=firstiter, finalizer="invalid")
+        except TypeError:
+            cur = sys.get_asyncgen_hooks()
+            self.assertIsNone(cur.firstiter)
+            self.assertIsNone(cur.finalizer)
+        else:
+            assert False
+
+        # gh-118473
+        try:
+            sys.set_asyncgen_hooks(firstiter="invalid", finalizer=finalizer)
+        except TypeError:
+            cur = sys.get_asyncgen_hooks()
+            self.assertIsNone(cur.firstiter)
+            self.assertIsNone(cur.finalizer)
+        else:
+            assert False
+
         sys.set_asyncgen_hooks(firstiter=firstiter)
         hooks = sys.get_asyncgen_hooks()
         self.assertIs(hooks.firstiter, firstiter)
         self.assertIs(hooks[0], firstiter)
         self.assertIs(hooks.finalizer, None)
         self.assertIs(hooks[1], None)
 
-        finalizer = lambda *a: None
         sys.set_asyncgen_hooks(finalizer=finalizer)
         hooks = sys.get_asyncgen_hooks()
         self.assertIs(hooks.firstiter, firstiter)

diff --git a/Misc/NEWS.d/next/Core and Builtins/2024-05-01-22-43-54.gh-issue-118473.QIvq9R.rst b/Misc/NEWS.d/next/Core and Builtins/2024-05-01-22-43-54.gh-issue-118473.QIvq9R.rst
@@ -0,0 +1 @@
+Fix set_asyncgen_hooks not to be partially set when raising TypeError
diff --git a/Python/ceval.c b/Python/ceval.c
@@ -2366,17 +2366,12 @@ _PyEval_GetCoroutineOriginTrackingDepth(void)
     return tstate->coroutine_origin_tracking_depth;
 }
 
-int
+void
 _PyEval_SetAsyncGenFirstiter(PyObject *firstiter)
 {
     PyThreadState *tstate = _PyThreadState_GET();
 
-    if (_PySys_Audit(tstate, "sys.set_asyncgen_hook_firstiter", NULL) < 0) {
-        return -1;
-    }
-
     Py_XSETREF(tstate->async_gen_firstiter, Py_XNewRef(firstiter));
-    return 0;
 }
 
 PyObject *
@@ -2386,17 +2381,12 @@ _PyEval_GetAsyncGenFirstiter(void)
     return tstate->async_gen_firstiter;
 }
 
-int
+void
 _PyEval_SetAsyncGenFinalizer(PyObject *finalizer)
 {
     PyThreadState *tstate = _PyThreadState_GET();
 
-    if (_PySys_Audit(tstate, "sys.set_asyncgen_hook_finalizer", NULL) < 0) {
-        return -1;
-    }
-
     Py_XSETREF(tstate->async_gen_finalizer, Py_XNewRef(finalizer));
-    return 0;
 }
 
 PyObject *

diff --git a/Python/sysmodule.c b/Python/sysmodule.c
@@ -1391,6 +1391,7 @@ static PyObject *
 sys_set_asyncgen_hooks(PyObject *self, PyObject *args, PyObject *kw)
 {
     static char *keywords[] = {"firstiter", "finalizer", NULL};
+    PyThreadState *tstate = NULL;
     PyObject *firstiter = NULL;
     PyObject *finalizer = NULL;
 
@@ -1400,34 +1401,45 @@ sys_set_asyncgen_hooks(PyObject *self, PyObject *args, PyObject *kw)
         return NULL;
     }
 
+    tstate = _PyThreadState_GET();
+
     if (finalizer && finalizer != Py_None) {
         if (!PyCallable_Check(finalizer)) {
             PyErr_Format(PyExc_TypeError,
                          "callable finalizer expected, got %.50s",
                          Py_TYPE(finalizer)->tp_name);
             return NULL;
         }
-        if (_PyEval_SetAsyncGenFinalizer(finalizer) < 0) {
+        if (_PySys_Audit(tstate, "sys.set_asyncgen_hook_finalizer", NULL) < 0) {
             return NULL;
         }
     }
-    else if (finalizer == Py_None && _PyEval_SetAsyncGenFinalizer(NULL) < 0) {
-        return NULL;
-    }
 
     if (firstiter && firstiter != Py_None) {
         if (!PyCallable_Check(firstiter)) {
             PyErr_Format(PyExc_TypeError,
                          "callable firstiter expected, got %.50s",
                          Py_TYPE(firstiter)->tp_name);
+
             return NULL;
         }
-        if (_PyEval_SetAsyncGenFirstiter(firstiter) < 0) {
+        if (_PySys_Audit(tstate, "sys.set_asyncgen_hook_firstiter", NULL) < 0) {
             return NULL;
         }
     }
-    else if (firstiter == Py_None && _PyEval_SetAsyncGenFirstiter(NULL) < 0) {
-        return NULL;
+
+    if (finalizer && finalizer != Py_None) {
+        _PyEval_SetAsyncGenFinalizer(finalizer);
+    }
+    else if (finalizer == Py_None) {
+        _PyEval_SetAsyncGenFinalizer(NULL);
+    }
+
+    if (firstiter && firstiter != Py_None) {
+        _PyEval_SetAsyncGenFirstiter(firstiter);
+    }
+    else if (firstiter == Py_None) {
+        _PyEval_SetAsyncGenFirstiter(NULL);
     }
 
     Py_RETURN_NONE;