https://github.com/python/cpython/commit/469d2e416c453b19d7a75fe31ceec732445e9ef2
commit: 469d2e416c453b19d7a75fe31ceec732445e9ef2
branch: main
author: Andrew Svetlov <[email protected]>
committer: asvetlov <[email protected]>
date: 2025-02-12T12:32:58+01:00
summary:
gh-129889: Support context manager protocol by contextvars.Token (#129888)
files:
A Misc/NEWS.d/next/Library/2025-02-10-09-45-49.gh-issue-129889.PBHXU5.rst
M Doc/library/contextvars.rst
M Doc/whatsnew/3.14.rst
M Lib/test/test_context.py
M Python/clinic/context.c.h
M Python/context.c
diff --git a/Doc/library/contextvars.rst b/Doc/library/contextvars.rst
index 2b1fb9fdd29cd8..3e3b30c724c631 100644
--- a/Doc/library/contextvars.rst
+++ b/Doc/library/contextvars.rst
@@ -101,6 +101,21 @@ Context Variables
the value of the variable to what it was before the corresponding
*set*.
+ The token supports :ref:`context manager protocol <context-managers>`
+ to restore the corresponding context variable value at the exit from
+ :keyword:`with` block::
+
+ var = ContextVar('var', default='default value')
+
+ with var.set('new value'):
+ assert var.get() == 'new value'
+
+ assert var.get() == 'default value'
+
+ .. versionadded:: next
+
+ Added support for usage as a context manager.
+
.. attribute:: Token.var
A read-only property. Points to the :class:`ContextVar` object
diff --git a/Doc/whatsnew/3.14.rst b/Doc/whatsnew/3.14.rst
index 3c7cc1b4529d32..e40b597ee52157 100644
--- a/Doc/whatsnew/3.14.rst
+++ b/Doc/whatsnew/3.14.rst
@@ -1,4 +1,3 @@
-
****************************
What's new in Python 3.14
****************************
@@ -362,6 +361,13 @@ concurrent.futures
supplying a *mp_context* to :class:`concurrent.futures.ProcessPoolExecutor`.
(Contributed by Gregory P. Smith in :gh:`84559`.)
+contextvars
+-----------
+
+* Support context manager protocol by :class:`contextvars.Token`.
+ (Contributed by Andrew Svetlov in :gh:`129889`.)
+
+
ctypes
------
diff --git a/Lib/test/test_context.py b/Lib/test/test_context.py
index 82d1797ab3b79e..f9cdcc3561e9d6 100644
--- a/Lib/test/test_context.py
+++ b/Lib/test/test_context.py
@@ -383,6 +383,115 @@ def sub(num):
tp.shutdown()
self.assertEqual(results, list(range(10)))
+ def test_token_contextmanager_with_default(self):
+ ctx = contextvars.Context()
+ c = contextvars.ContextVar('c', default=42)
+
+ def fun():
+ with c.set(36):
+ self.assertEqual(c.get(), 36)
+
+ self.assertEqual(c.get(), 42)
+
+ ctx.run(fun)
+
+ def test_token_contextmanager_without_default(self):
+ ctx = contextvars.Context()
+ c = contextvars.ContextVar('c')
+
+ def fun():
+ with c.set(36):
+ self.assertEqual(c.get(), 36)
+
+ with self.assertRaisesRegex(LookupError, "<ContextVar name='c'"):
+ c.get()
+
+ ctx.run(fun)
+
+ def test_token_contextmanager_on_exception(self):
+ ctx = contextvars.Context()
+ c = contextvars.ContextVar('c', default=42)
+
+ def fun():
+ with c.set(36):
+ self.assertEqual(c.get(), 36)
+ raise ValueError("custom exception")
+
+ self.assertEqual(c.get(), 42)
+
+ with self.assertRaisesRegex(ValueError, "custom exception"):
+ ctx.run(fun)
+
+ def test_token_contextmanager_reentrant(self):
+ ctx = contextvars.Context()
+ c = contextvars.ContextVar('c', default=42)
+
+ def fun():
+ token = c.set(36)
+ with self.assertRaisesRegex(
+ RuntimeError,
+ "<Token .+ has already been used once"
+ ):
+ with token:
+ with token:
+ self.assertEqual(c.get(), 36)
+
+ self.assertEqual(c.get(), 42)
+
+ ctx.run(fun)
+
+ def test_token_contextmanager_multiple_c_set(self):
+ ctx = contextvars.Context()
+ c = contextvars.ContextVar('c', default=42)
+
+ def fun():
+ with c.set(36):
+ self.assertEqual(c.get(), 36)
+ c.set(24)
+ self.assertEqual(c.get(), 24)
+ c.set(12)
+ self.assertEqual(c.get(), 12)
+
+ self.assertEqual(c.get(), 42)
+
+ ctx.run(fun)
+
+ def test_token_contextmanager_with_explicit_reset_the_same_token(self):
+ ctx = contextvars.Context()
+ c = contextvars.ContextVar('c', default=42)
+
+ def fun():
+ with self.assertRaisesRegex(
+ RuntimeError,
+ "<Token .+ has already been used once"
+ ):
+ with c.set(36) as token:
+ self.assertEqual(c.get(), 36)
+ c.reset(token)
+
+ self.assertEqual(c.get(), 42)
+
+ self.assertEqual(c.get(), 42)
+
+ ctx.run(fun)
+
+ def test_token_contextmanager_with_explicit_reset_another_token(self):
+ ctx = contextvars.Context()
+ c = contextvars.ContextVar('c', default=42)
+
+ def fun():
+ with c.set(36):
+ self.assertEqual(c.get(), 36)
+
+ token = c.set(24)
+ self.assertEqual(c.get(), 24)
+ c.reset(token)
+ self.assertEqual(c.get(), 36)
+
+ self.assertEqual(c.get(), 42)
+
+ ctx.run(fun)
+
# HAMT Tests
diff --git
a/Misc/NEWS.d/next/Library/2025-02-10-09-45-49.gh-issue-129889.PBHXU5.rst
b/Misc/NEWS.d/next/Library/2025-02-10-09-45-49.gh-issue-129889.PBHXU5.rst
new file mode 100644
index 00000000000000..f0880e5de8412e
--- /dev/null
+++ b/Misc/NEWS.d/next/Library/2025-02-10-09-45-49.gh-issue-129889.PBHXU5.rst
@@ -0,0 +1,2 @@
+Support context manager protocol by :class:`contextvars.Token`. Patch by
+Andrew Svetlov.
diff --git a/Python/clinic/context.c.h b/Python/clinic/context.c.h
index 71f05aa02a51e7..0adde76d7c3cb1 100644
--- a/Python/clinic/context.c.h
+++ b/Python/clinic/context.c.h
@@ -179,4 +179,55 @@ PyDoc_STRVAR(_contextvars_ContextVar_reset__doc__,
#define _CONTEXTVARS_CONTEXTVAR_RESET_METHODDEF \
{"reset", (PyCFunction)_contextvars_ContextVar_reset, METH_O,
_contextvars_ContextVar_reset__doc__},
-/*[clinic end generated code: output=444567eaf0df25e0 input=a9049054013a1b77]*/
+
+PyDoc_STRVAR(token_enter__doc__,
+"__enter__($self, /)\n"
+"--\n"
+"\n"
+"Enter into Token context manager.");
+
+#define TOKEN_ENTER_METHODDEF \
+ {"__enter__", (PyCFunction)token_enter, METH_NOARGS, token_enter__doc__},
+
+static PyObject *
+token_enter_impl(PyContextToken *self);
+
+static PyObject *
+token_enter(PyObject *self, PyObject *Py_UNUSED(ignored))
+{
+ return token_enter_impl((PyContextToken *)self);
+}
+
+PyDoc_STRVAR(token_exit__doc__,
+"__exit__($self, type, val, tb, /)\n"
+"--\n"
+"\n"
+"Exit from Token context manager, restore the linked ContextVar.");
+
+#define TOKEN_EXIT_METHODDEF \
+ {"__exit__", _PyCFunction_CAST(token_exit), METH_FASTCALL,
token_exit__doc__},
+
+static PyObject *
+token_exit_impl(PyContextToken *self, PyObject *type, PyObject *val,
+ PyObject *tb);
+
+static PyObject *
+token_exit(PyObject *self, PyObject *const *args, Py_ssize_t nargs)
+{
+ PyObject *return_value = NULL;
+ PyObject *type;
+ PyObject *val;
+ PyObject *tb;
+
+ if (!_PyArg_CheckPositional("__exit__", nargs, 3, 3)) {
+ goto exit;
+ }
+ type = args[0];
+ val = args[1];
+ tb = args[2];
+ return_value = token_exit_impl((PyContextToken *)self, type, val, tb);
+
+exit:
+ return return_value;
+}
+/*[clinic end generated code: output=01987cdbf68a951a input=a9049054013a1b77]*/
diff --git a/Python/context.c b/Python/context.c
index bb1aa42b9c5e4f..dfdde7d1fa723f 100644
--- a/Python/context.c
+++ b/Python/context.c
@@ -1231,9 +1231,47 @@ static PyGetSetDef PyContextTokenType_getsetlist[] = {
{NULL}
};
+/*[clinic input]
+_contextvars.Token.__enter__ as token_enter
+
+Enter into Token context manager.
+[clinic start generated code]*/
+
+static PyObject *
+token_enter_impl(PyContextToken *self)
+/*[clinic end generated code: output=9af4d2054e93fb75 input=41a3d6c4195fd47a]*/
+{
+ return Py_NewRef(self);
+}
+
+/*[clinic input]
+_contextvars.Token.__exit__ as token_exit
+
+ type: object
+ val: object
+ tb: object
+ /
+
+Exit from Token context manager, restore the linked ContextVar.
+[clinic start generated code]*/
+
+static PyObject *
+token_exit_impl(PyContextToken *self, PyObject *type, PyObject *val,
+ PyObject *tb)
+/*[clinic end generated code: output=3e6a1c95d3da703a input=7f117445f0ccd92e]*/
+{
+ int ret = PyContextVar_Reset((PyObject *)self->tok_var, (PyObject *)self);
+ if (ret < 0) {
+ return NULL;
+ }
+ Py_RETURN_NONE;
+}
+
static PyMethodDef PyContextTokenType_methods[] = {
{"__class_getitem__", Py_GenericAlias,
METH_O|METH_CLASS, PyDoc_STR("See PEP 585")},
+ TOKEN_ENTER_METHODDEF
+ TOKEN_EXIT_METHODDEF
{NULL}
};
_______________________________________________
Python-checkins mailing list -- [email protected]
To unsubscribe send an email to [email protected]
https://mail.python.org/mailman3/lists/python-checkins.python.org/
Member address: [email protected]
