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 ` + 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, "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} };