gh-127124: Change context watcher callback to a callable object

Doc/c-api/contextvars.rst

@@ -101,21 +101,76 @@ Context object management functions:
    current context for the current thread.  Returns ``0`` on success,
    and ``-1`` on error.
 
-.. c:function:: int PyContext_AddWatcher(PyContext_WatchCallback callback)
-
-   Register *callback* as a context object watcher for the current interpreter.
-   Return an ID which may be passed to :c:func:`PyContext_ClearWatcher`.
-   In case of error (e.g. no more watcher IDs available),
-   return ``-1`` and set an exception.
+.. c:function:: int PyContext_AddWatcher(PyObject *callback)
+
+   Registers the callable object *callback* as a context object watcher for the
+   current interpreter.  When a context event occurs, *callback* is called with
+   two arguments:
+
+   #. An event type ID from :c:type:`PyContextEvent`.
+   #. An object containing event-specific supplemental data; see
+      :c:type:`PyContextEvent` for details.
+
+   Any exception raised by *callback* will be printed as an unraisable
+   exception as if by a call to :c:func:`PyErr_FormatUnraisable`, then
+   discarded.
+
+   On success, this function returns a non-negative ID which may be passed to
+   :c:func:`PyContext_ClearWatcher` to unregister the callback and remove the
+   reference this function adds to *callback*.  Sets an exception and returns
+   ``-1`` on error (e.g., no more watcher IDs available).
+
+   Example using a C function as the callback::
+
+      static PyObject *
+      my_callback(PyObject *self, PyObject *const *args, Py_ssize_t nargs)
+      {
+          if (PyVectorcall_NARGS(nargs) != 2) {
+              PyErr_Format(PyExc_TypeError, "want 2 args, got %zd", nargs);
+              return NULL;
+          }
+          int event = PyLong_AsInt(args[0]);
+          if (event == -1 && PyErr_Occurred()) {
+              return NULL;
+          }
+          if (event != Py_CONTEXT_SWITCHED) {
+              Py_RETURN_NONE;
+          }
+          PyObject *ctx = args[1];
+
+          // Do something interesting with self and ctx here.
+
+          Py_RETURN_NONE;
+      }
+
+      PyMethodDef my_callback_md = {
+          .ml_name = "my_callback",
+          .ml_meth = (PyCFunction)(void *)&my_callback,
+          .ml_flags = METH_FASTCALL,
+          .ml_doc = NULL,
+      };
+
+      int
+      register_my_callback(PyObject *callback_state)
+      {
+          PyObject *cb = PyCFunction_New(&my_callback_md, callback_state);
+          if (cb == NULL) {
+              return -1;
+          }
+          int id = PyContext_AddWatcher(cb);
+          Py_CLEAR(cb);
+          return id;
+      }
 
    .. versionadded:: 3.14
 
 .. c:function:: int PyContext_ClearWatcher(int watcher_id)
 
-   Clear watcher identified by *watcher_id* previously returned from
-   :c:func:`PyContext_AddWatcher` for the current interpreter.
-   Return ``0`` on success, or ``-1`` and set an exception on error
-   (e.g. if the given *watcher_id* was never registered.)
+   Clears the watcher identified by *watcher_id* previously returned from
+   :c:func:`PyContext_AddWatcher` for the current interpreter, and removes the
+   reference created for the registered callback object.  Returns ``0`` on
+   success, or sets an exception and returns ``-1`` on error (e.g., if the
+   given *watcher_id* was never registered).
 
    .. versionadded:: 3.14
 
@@ -130,23 +185,6 @@ Context object management functions:
 
    .. versionadded:: 3.14
 
-.. c:type:: int (*PyContext_WatchCallback)(PyContextEvent event, PyObject *obj)
-
-   Context object watcher callback function.  The object passed to the callback
-   is event-specific; see :c:type:`PyContextEvent` for details.
-
-   If the callback returns with an exception set, it must return ``-1``; this
-   exception will be printed as an unraisable exception using
-   :c:func:`PyErr_FormatUnraisable`. Otherwise it should return ``0``.
-
-   There may already be a pending exception set on entry to the callback. In
-   this case, the callback should return ``0`` with the same exception still
-   set. This means the callback may not call any other API that can set an
-   exception unless it saves and clears the exception state first, and restores
-   it before returning.
-
-   .. versionadded:: 3.14
-
 
 Context variable functions:
 

Include/cpython/context.h

@@ -28,37 +28,10 @@ PyAPI_FUNC(int) PyContext_Enter(PyObject *);
 PyAPI_FUNC(int) PyContext_Exit(PyObject *);
 
 typedef enum {
-    /*
-     * The current context has switched to a different context.  The object
-     * passed to the watch callback is the now-current contextvars.Context
-     * object, or None if no context is current.
-     */
     Py_CONTEXT_SWITCHED = 1,
 } PyContextEvent;
 
-/*
- * Context object watcher callback function.  The object passed to the callback
- * is event-specific; see PyContextEvent for details.
- *
- * if the callback returns with an exception set, it must return -1. Otherwise
- * it should return 0
- */
-typedef int (*PyContext_WatchCallback)(PyContextEvent, PyObject *);
-
-/*
- * Register a per-interpreter callback that will be invoked for context object
- * enter/exit events.
- *
- * Returns a handle that may be passed to PyContext_ClearWatcher on success,
- * or -1 and sets and error if no more handles are available.
- */
-PyAPI_FUNC(int) PyContext_AddWatcher(PyContext_WatchCallback callback);
-
-/*
- * Clear the watcher associated with the watcher_id handle.
- *
- * Returns 0 on success or -1 if no watcher exists for the provided id.
- */
+PyAPI_FUNC(int) PyContext_AddWatcher(PyObject *callback);
 PyAPI_FUNC(int) PyContext_ClearWatcher(int watcher_id);
 
 /* Create a new context variable.

Include/internal/pycore_interp.h

@@ -242,7 +242,7 @@ struct _is {
     PyObject *audit_hooks;
     PyType_WatchCallback type_watchers[TYPE_MAX_WATCHERS];
     PyCode_WatchCallback code_watchers[CODE_MAX_WATCHERS];
-    PyContext_WatchCallback context_watchers[CONTEXT_MAX_WATCHERS];
+    PyObject *context_watchers[CONTEXT_MAX_WATCHERS];
     // One bit is set for each non-NULL entry in code_watchers
     uint8_t active_code_watchers;
     uint8_t active_context_watchers;

Lib/test/test_capi/test_watchers.py

@@ -1,3 +1,4 @@
+import contextlib
 import unittest
 import contextvars
 
@@ -589,60 +590,49 @@ def test_allocate_too_many_watchers(self):
 
 class TestContextObjectWatchers(unittest.TestCase):
     @contextmanager
-    def context_watcher(self, which_watcher):
-        wid = _testcapi.add_context_watcher(which_watcher)
+    def context_watcher(self, cb=None):
+        log = None
+        if cb is None:
+            log = []
+            def cb(event, ctx):
+                self.assertEqual(event, _testcapi.Py_CONTEXT_SWITCHED)
+                log.append(ctx)
+        wid = _testcapi.add_context_watcher(cb)
         try:
-            switches = _testcapi.get_context_switches(which_watcher)
-        except ValueError:
-            switches = None
-        try:
-            yield switches
+            yield log
         finally:
             _testcapi.clear_context_watcher(wid)
 
-    def assert_event_counts(self, want_0, want_1):
-        self.assertEqual(len(_testcapi.get_context_switches(0)), want_0)
-        self.assertEqual(len(_testcapi.get_context_switches(1)), want_1)
-
     def test_context_object_events_dispatched(self):
-        # verify that all counts are zero before any watchers are registered
-        self.assert_event_counts(0, 0)
-
-        # verify that all counts remain zero when a context object is
-        # entered and exited with no watchers registered
         ctx = contextvars.copy_context()
-        ctx.run(self.assert_event_counts, 0, 0)
-        self.assert_event_counts(0, 0)
-
-        # verify counts are as expected when first watcher is registered
-        with self.context_watcher(0):
-            self.assert_event_counts(0, 0)
-            ctx.run(self.assert_event_counts, 1, 0)
-            self.assert_event_counts(2, 0)
-
-            # again with second watcher registered
-            with self.context_watcher(1):
-                self.assert_event_counts(2, 0)
-                ctx.run(self.assert_event_counts, 3, 1)
-                self.assert_event_counts(4, 2)
-
-        # verify counts are reset and don't change after both watchers are cleared
-        ctx.run(self.assert_event_counts, 0, 0)
-        self.assert_event_counts(0, 0)
+        with self.context_watcher() as switches_0:
+            self.assertEqual(len(switches_0), 0)
+            ctx.run(lambda: self.assertEqual(len(switches_0), 1))
+            self.assertEqual(len(switches_0), 2)
+            with self.context_watcher() as switches_1:
+                self.assertEqual((len(switches_0), len(switches_1)), (2, 0))
+                ctx.run(lambda: self.assertEqual(
+                    (len(switches_0), len(switches_1)), (3, 1)))
+                self.assertEqual((len(switches_0), len(switches_1)), (4, 2))
 
     def test_callback_error(self):
         ctx_outer = contextvars.copy_context()
         ctx_inner = contextvars.copy_context()
         unraisables = []
 
+        def _cb(event, ctx):
+            raise RuntimeError('boom!')
+
         def _in_outer():
-            with self.context_watcher(2):
+            with self.context_watcher(_cb):
                 with catch_unraisable_exception() as cm:
                     ctx_inner.run(lambda: unraisables.append(cm.unraisable))
                     unraisables.append(cm.unraisable)
 
         try:
             ctx_outer.run(_in_outer)
+            self.assertEqual([x is not None for x in unraisables],
+                             [True, True])
             self.assertEqual([x.err_msg for x in unraisables],
                              ["Exception ignored in Py_CONTEXT_SWITCHED "
                               f"watcher callback for {ctx!r}"
@@ -656,21 +646,24 @@ def _in_outer():
     def test_clear_out_of_range_watcher_id(self):
         with self.assertRaisesRegex(ValueError, r"Invalid context watcher ID -1"):
             _testcapi.clear_context_watcher(-1)
-        with self.assertRaisesRegex(ValueError, r"Invalid context watcher ID 8"):
-            _testcapi.clear_context_watcher(8)  # CONTEXT_MAX_WATCHERS = 8
+        with self.assertRaisesRegex(ValueError, f"Invalid context watcher ID {_testcapi.CONTEXT_MAX_WATCHERS}"):
+            _testcapi.clear_context_watcher(_testcapi.CONTEXT_MAX_WATCHERS)
 
     def test_clear_unassigned_watcher_id(self):
         with self.assertRaisesRegex(ValueError, r"No context watcher set for ID 1"):
             _testcapi.clear_context_watcher(1)
 
     def test_allocate_too_many_watchers(self):
-        with self.assertRaisesRegex(RuntimeError, r"no more context watcher IDs available"):
-            _testcapi.allocate_too_many_context_watchers()
+        with contextlib.ExitStack() as stack:
+            for i in range(_testcapi.CONTEXT_MAX_WATCHERS):
+                stack.enter_context(self.context_watcher())
+            with self.assertRaisesRegex(RuntimeError, r"no more context watcher IDs available"):
+                stack.enter_context(self.context_watcher())
 
     def test_exit_base_context(self):
         ctx = contextvars.Context()
         _testcapi.clear_context_stack()
-        with self.context_watcher(0) as switches:
+        with self.context_watcher() as switches:
             ctx.run(lambda: None)
         self.assertEqual(switches, [ctx, None])
 

Misc/NEWS.d/next/C_API/2024-11-22-18-38-33.gh-issue-127124.6k6Qj7.rst

@@ -0,0 +1,2 @@
+Changed :c:func:`PyContext_AddWatcher` to take a callable object instead of a C
+function pointer so that the callback can have non-global state.