Add PyContextVar_GetChanged() method.

Author: nascheme

Doc/c-api/contextvars.rst

@@ -182,3 +182,27 @@ Context variable functions:
    Reset the state of the *var* context variable to that it was in before
    :c:func:`PyContextVar_Set` that returned the *token* was called.
    This function returns ``0`` on success and ``-1`` on error.
+
+.. c:function:: int PyContextVar_GetChanged(PyObject *var, PyObject *default_value, PyObject **value, int *changed)
+
+   Like :c:func:`PyContextVar_Get`, but also reports whether the variable was
+   changed in the current context scope.  This combines a value lookup with a
+   change check in a single HAMT lookup.
+
+   Returns ``-1`` if an error has occurred during lookup, and ``0`` if no
+   error occurred, whether or not a value was found.
+
+   On success, *\*value* is set following the same rules as
+   :c:func:`PyContextVar_Get`.  *\*changed* is set to ``1`` if the variable
+   was changed (via :c:func:`PyContextVar_Set`) in the current context scope
+   (i.e. within the current :meth:`~contextvars.Context.run` call) with a
+   value that is a different object than the inherited one.  Otherwise
+   *\*changed* is set to ``0``.  If the value was not found, *\*changed* is
+   always ``0``.
+
+   If the current context was never entered (no :meth:`~contextvars.Context.run`
+   is active), all existing bindings are considered "changed".
+
+   Except for ``NULL``, the function returns a new reference via *\*value*.
+
+   .. versionadded:: 3.15

Doc/data/refcounts.dat

@@ -400,6 +400,12 @@ PyContextVar_Reset:int:::
 PyContextVar_Reset:PyObject*:var:0:
 PyContextVar_Reset:PyObject*:token:-1:
 
+PyContextVar_GetChanged:int:::
+PyContextVar_GetChanged:PyObject*:var:0:
+PyContextVar_GetChanged:PyObject*:default_value:0:
+PyContextVar_GetChanged:PyObject**:value:+1:???
+PyContextVar_GetChanged:int*:changed::
+
 PyCFunction_New:PyObject*::+1:
 PyCFunction_New:PyMethodDef*:ml::
 PyCFunction_New:PyObject*:self:+1:

Doc/library/contextvars.rst

@@ -119,6 +119,41 @@ Context Variables
 
       The same *token* cannot be used twice.
 
+   .. method:: get_changed([default])
+
+      Like :meth:`ContextVar.get`, but returns a tuple ``(value, changed)``
+      where *changed* indicates whether the variable was changed in the
+      current context scope.
+
+      A variable is considered "changed" if :meth:`ContextVar.set` has been
+      called on it within the current :meth:`Context.run` call with a value
+      that is a different object than the inherited one.  Variables inherited
+      unchanged from a parent context scope are not considered "changed".
+      If no :meth:`Context.run` is active, all existing bindings are
+      considered "changed".  When the value comes from a default, *changed*
+      is always ``False``.
+
+      This is useful when a context variable holds a mutable object that
+      needs to be copied on first access in a new context scope to ensure
+      modifications are local to that scope::
+
+         _ctx_var = ContextVar('ctx_var')
+
+         def get_ctx():
+             try:
+                 ctx, changed = _ctx_var.get_changed()
+             except LookupError:
+                 ctx = default_context()
+                 _ctx_var.set(ctx)
+                 return ctx
+
+             if not changed:
+                 ctx = ctx.copy()
+                 _ctx_var.set(ctx)
+             return ctx
+
+      .. versionadded:: 3.15
+
 
 .. class:: Token
 

Include/cpython/context.h

@@ -100,6 +100,26 @@ PyAPI_FUNC(PyObject *) PyContextVar_Set(PyObject *var, PyObject *value);
 PyAPI_FUNC(int) PyContextVar_Reset(PyObject *var, PyObject *token);
 
 
+/* Get a value for the variable and check if it was changed.
+
+   Like PyContextVar_Get, but also reports whether the variable was
+   changed in the current context scope via a single HAMT lookup.
+
+   Returns -1 if an error occurred during lookup.
+
+   Returns 0 if no error occurred.  In this case:
+
+   - *value will be set the same as for PyContextVar_Get.
+   - *changed will be set to 1 if the variable was changed in the
+     current context scope, 0 otherwise.  If the variable was not
+     found, *changed is always 0.
+
+   '*value' will be a new ref, if not NULL.
+*/
+PyAPI_FUNC(int) PyContextVar_GetChanged(
+    PyObject *var, PyObject *default_value, PyObject **value, int *changed);
+
+
 #ifdef __cplusplus
 }
 #endif

Include/internal/pycore_context.h

@@ -24,6 +24,7 @@ struct _pycontextobject {
     PyObject_HEAD
     PyContext *ctx_prev;
     PyHamtObject *ctx_vars;
+    PyHamtObject *ctx_vars_origin;  /* snapshot of ctx_vars at Enter time */
     PyObject *ctx_weakreflist;
     int ctx_entered;
 };

Lib/test/test_context.py

@@ -586,6 +586,226 @@ def __eq__(self, other):
         ctx2.run(var.set, ReentrantHash())
         ctx1 == ctx2
 
+    def test_get_changed_outside_run(self):
+        # Outside any Context.run(), bindings are considered "changed"
+        v = contextvars.ContextVar('v', default='dflt')
+        val, changed = v.get_changed()
+        self.assertEqual(val, 'dflt')
+        self.assertFalse(changed)  # default value, not changed
+        v.set(42)
+        val, changed = v.get_changed()
+        self.assertEqual(val, 42)
+        self.assertTrue(changed)  # set in base context
+
+    def test_get_changed_inherited(self):
+        # Inherited bindings are not considered "changed"
+        v = contextvars.ContextVar('v')
+        v.set('parent')
+        ctx = contextvars.copy_context()
+
+        def check():
+            val, changed = v.get_changed()
+            self.assertEqual(val, 'parent')
+            self.assertFalse(changed)
+        ctx.run(check)
+
+    def test_get_changed_after_set(self):
+        # After set() inside Context.run(), changed is True
+        v = contextvars.ContextVar('v')
+        v.set('parent')
+        ctx = contextvars.copy_context()
+
+        def check():
+            val, changed = v.get_changed()
+            self.assertFalse(changed)
+            v.set('child')
+            val, changed = v.get_changed()
+            self.assertEqual(val, 'child')
+            self.assertTrue(changed)
+        ctx.run(check)
+
+    def test_get_changed_new_var_in_run(self):
+        # A variable set for the first time inside run() is "changed"
+        v = contextvars.ContextVar('v')
+        ctx = contextvars.copy_context()
+
+        def check():
+            with self.assertRaises(LookupError):
+                v.get_changed()
+            v.set('new')
+            val, changed = v.get_changed()
+            self.assertEqual(val, 'new')
+            self.assertTrue(changed)
+        ctx.run(check)
+
+    def test_get_changed_not_set_with_default(self):
+        # A variable not set but with default: changed is False
+        v = contextvars.ContextVar('v', default='dflt')
+        ctx = contextvars.copy_context()
+
+        def check():
+            val, changed = v.get_changed()
+            self.assertEqual(val, 'dflt')
+            self.assertFalse(changed)
+        ctx.run(check)
+
+    def test_get_changed_not_set_no_default(self):
+        # A variable that has never been set and has no default
+        v = contextvars.ContextVar('v')
+        ctx = contextvars.copy_context()
+
+        def check():
+            with self.assertRaises(LookupError):
+                v.get_changed()
+        ctx.run(check)
+
+    def test_get_changed_explicit_default_arg(self):
+        # Passing a default argument to get_changed()
+        v = contextvars.ContextVar('v')
+        ctx = contextvars.copy_context()
+
+        def check():
+            val, changed = v.get_changed('fallback')
+            self.assertEqual(val, 'fallback')
+            self.assertFalse(changed)
+        ctx.run(check)
+
+    def test_get_changed_set_same_object(self):
+        # Setting to the exact same object does not count as "changed"
+        # because the HAMT recognizes the identical key-value pair
+        obj = object()
+        v = contextvars.ContextVar('v')
+        v.set(obj)
+        ctx = contextvars.copy_context()
+
+        def check():
+            val, changed = v.get_changed()
+            self.assertIs(val, obj)
+            self.assertFalse(changed)
+            v.set(obj)  # same object
+            val, changed = v.get_changed()
+            self.assertIs(val, obj)
+            self.assertFalse(changed)
+        ctx.run(check)
+
+    def test_get_changed_set_different_object(self):
+        # Setting to a different object counts as "changed"
+        v = contextvars.ContextVar('v')
+        v.set([1, 2, 3])
+        ctx = contextvars.copy_context()
+
+        def check():
+            val, changed = v.get_changed()
+            self.assertFalse(changed)
+            v.set([1, 2, 3])  # equal value, different object
+            val, changed = v.get_changed()
+            self.assertTrue(changed)
+        ctx.run(check)
+
+    def test_get_changed_after_reset(self):
+        # After reset(), the variable reverts to its inherited state
+        v = contextvars.ContextVar('v')
+        v.set('original')
+        ctx = contextvars.copy_context()
+
+        def check():
+            val, changed = v.get_changed()
+            self.assertFalse(changed)
+            tok = v.set('modified')
+            val, changed = v.get_changed()
+            self.assertTrue(changed)
+            v.reset(tok)
+            val, changed = v.get_changed()
+            self.assertFalse(changed)
+        ctx.run(check)
+
+    def test_get_changed_multiple_vars(self):
+        # Changing one variable does not affect get_changed() for others
+        v1 = contextvars.ContextVar('v1')
+        v2 = contextvars.ContextVar('v2')
+        v1.set('a')
+        v2.set('b')
+        ctx = contextvars.copy_context()
+
+        def check():
+            _, changed1 = v1.get_changed()
+            _, changed2 = v2.get_changed()
+            self.assertFalse(changed1)
+            self.assertFalse(changed2)
+            v1.set('a2')
+            _, changed1 = v1.get_changed()
+            _, changed2 = v2.get_changed()
+            self.assertTrue(changed1)
+            self.assertFalse(changed2)
+        ctx.run(check)
+
+    def test_get_changed_nested_run(self):
+        # get_changed() reflects the innermost Context.run() scope
+        v = contextvars.ContextVar('v')
+        v.set('root')
+        ctx1 = contextvars.copy_context()
+
+        def outer():
+            _, changed = v.get_changed()
+            self.assertFalse(changed)
+            v.set('outer')
+            _, changed = v.get_changed()
+            self.assertTrue(changed)
+            ctx2 = contextvars.copy_context()
+
+            def inner():
+                # inherited 'outer' from ctx1, not changed in ctx2
+                val, changed = v.get_changed()
+                self.assertEqual(val, 'outer')
+                self.assertFalse(changed)
+                v.set('inner')
+                val, changed = v.get_changed()
+                self.assertEqual(val, 'inner')
+                self.assertTrue(changed)
+            ctx2.run(inner)
+
+            # after inner run exits, outer's state is restored
+            _, changed = v.get_changed()
+            self.assertTrue(changed)
+        ctx1.run(outer)
+
+    def test_get_changed_with_threads(self):
+        # get_changed() works correctly in a thread with copied context
+        import threading
+        v = contextvars.ContextVar('v')
+        v.set('parent')
+        ctx = contextvars.copy_context()
+        results = {}
+
+        def thread_func():
+            val, changed = v.get_changed()
+            results['inherited'] = changed
+            results['value'] = val
+            v.set('thread')
+            val, changed = v.get_changed()
+            results['after_set'] = changed
+
+        t = threading.Thread(target=ctx.run, args=(thread_func,))
+        t.start()
+        t.join()
+        self.assertFalse(results['inherited'])
+        self.assertEqual(results['value'], 'parent')
+        self.assertTrue(results['after_set'])
+
+    def test_get_changed_empty_context_run(self):
+        # Running in a brand new empty context
+        v = contextvars.ContextVar('v')
+        ctx = contextvars.Context()
+
+        def check():
+            with self.assertRaises(LookupError):
+                v.get_changed()
+            v.set('value')
+            val, changed = v.get_changed()
+            self.assertEqual(val, 'value')
+            self.assertTrue(changed)
+        ctx.run(check)
+
 
 # HAMT Tests