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

This enables developers to associate state with the callback without
relying on globals.

Also, refactor the tests for improved readability and extensibility,
and to cover the new state object.

Author: rhansen

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

@@ -36,29 +36,7 @@ typedef enum {
     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.