gh-127124: Pass optional state to context watcher callback

Doc/c-api/contextvars.rst

@@ -101,21 +101,25 @@ 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)
+.. c:function:: int PyContext_AddWatcher(PyContext_WatchCallback *callback, PyObject *cbarg)
 
- 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.
+ Registers *callback* as a context object watcher for the current
+ interpreter, and *cbarg* (which may be NULL; if not, this function creates a
+ new reference) as the object to pass as the callback's first parameter. On
+ success, returns a non-negative ID which may be passed to
+ :c:func:`PyContext_ClearWatcher` to unregister the callback and remove the
+ added reference to *cbarg*. Sets an exception and returns ``-1`` on error
+ (e.g., no more watcher IDs available).
 
  .. 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 *cbarg* object that was registered with the
+ callback. 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,10 +134,12 @@ Context object management functions:
 
  .. versionadded:: 3.14
 
-.. c:type:: int (*PyContext_WatchCallback)(PyContextEvent event, PyObject *obj)
+.. c:type:: int PyContext_WatchCallback(PyObject *cbarg, PyContextEvent event, PyObject *obj)
 
- Context object watcher callback function. The object passed to the callback
- is event-specific; see :c:type:`PyContextEvent` for details.
+ Context object watcher callback function. *cbarg* is the same object
+ registered in the call to :c:func:`PyContext_AddWatcher`, as a borrowed
+ reference if non-NULL. The *obj* object 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

Include/cpython/context.h

@@ -28,37 +28,13 @@ 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.
- */
+typedef int PyContext_WatchCallback(
+ PyObject *cbarg, PyContextEvent event, PyObject *obj);
+PyAPI_FUNC(int) PyContext_AddWatcher(
+ PyContext_WatchCallback *callback, PyObject *cbarg);
 PyAPI_FUNC(int) PyContext_ClearWatcher(int watcher_id);
 
 /* Create a new context variable.

Include/internal/pycore_interp.h

@@ -242,7 +242,10 @@ 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];
+ struct{
+ PyContext_WatchCallback *callback;
+ PyObject *arg;
+ } 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,43 @@ 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, cfg=None):
+ if cfg is None:
+ cfg ={}
+ cfg.setdefault('log', [])
+ wid = _testcapi.add_context_watcher(cfg)
  try:
- switches = _testcapi.get_context_switches(which_watcher)
- except ValueError:
- switches = None
- try:
- yield switches
+ yield cfg.get('log', None)
  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 _in_outer():
- with self.context_watcher(2):
+ with self.context_watcher(cfg={'err': RuntimeError('boom!')}):
  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 +640,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 @@
+Added optional callback state to :c:func:`PyContext_AddWatcher` and
+:c:type:`PyContext_WatchCallback`.