From c214680432cc5b101760de821592593fc1915d72 Mon Sep 17 00:00:00 2001 From: Andrew Svetlov Date: Fri, 3 Jan 2025 21:00:01 +0100 Subject: [PATCH 1/8] Implement context manager for Token --- Lib/test/test_context.py | 25 +++++++++++++++++++++++++ Python/context.c | 20 ++++++++++++++++++++ 2 files changed, 45 insertions(+) diff --git a/Lib/test/test_context.py b/Lib/test/test_context.py index 82d1797ab3b79e..33765f45124f9f 100644 --- a/Lib/test/test_context.py +++ b/Lib/test/test_context.py @@ -383,6 +383,31 @@ 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")}, + + {"__enter__", (PyCFunction)token_enter, METH_NOARGS, NULL}, + {"__exit__", (PyCFunction)token_exit, METH_VARARGS, NULL}, + {NULL} }; From 0400e8e6145f804cb23e0d9b2558337a173d834d Mon Sep 17 00:00:00 2001 From: Andrew Svetlov Date: Sun, 9 Feb 2025 11:27:24 +0100 Subject: [PATCH 2/8] Use clinic for stubs generation --- Python/clinic/context.c.h | 53 ++++++++++++++++++++++++++++++++++++++- Python/context.c | 30 +++++++++++++++++----- 2 files changed, 76 insertions(+), 7 deletions(-) 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 ea25af7e23b823..4b2050cb51aa58 100644 --- a/Python/context.c +++ b/Python/context.c @@ -1231,14 +1231,34 @@ 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(PyContextToken *self, PyObject *args) +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(PyContextToken *self, PyObject *args) +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) { @@ -1250,10 +1270,8 @@ token_exit(PyContextToken *self, PyObject *args) static PyMethodDef PyContextTokenType_methods[] = { {"__class_getitem__", Py_GenericAlias, METH_O|METH_CLASS, PyDoc_STR("See PEP 585")}, - - {"__enter__", (PyCFunction)token_enter, METH_NOARGS, NULL}, - {"__exit__", (PyCFunction)token_exit, METH_VARARGS, NULL}, - + BLOB_ENTER_METHODDEF + BLOB_EXIT_METHODDEF {NULL} }; From 61abb23dfb61103b08cb9939daf2091045ce0595 Mon Sep 17 00:00:00 2001 From: Andrew Svetlov Date: Sun, 9 Feb 2025 12:28:51 +0100 Subject: [PATCH 3/8] Update docs --- Doc/library/contextvars.rst | 15 +++++++++++++++ 1 file changed, 15 insertions(+) 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 From 191f7a14e6ba27022425e9c4367f364833157b8e Mon Sep 17 00:00:00 2001 From: Andrew Svetlov Date: Sun, 9 Feb 2025 12:39:44 +0100 Subject: [PATCH 4/8] fix --- Python/context.c | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/Python/context.c b/Python/context.c index 4b2050cb51aa58..dfdde7d1fa723f 100644 --- a/Python/context.c +++ b/Python/context.c @@ -1270,8 +1270,8 @@ token_exit_impl(PyContextToken *self, PyObject *type, PyObject *val, static PyMethodDef PyContextTokenType_methods[] = { {"__class_getitem__", Py_GenericAlias, METH_O|METH_CLASS, PyDoc_STR("See PEP 585")}, - BLOB_ENTER_METHODDEF - BLOB_EXIT_METHODDEF + TOKEN_ENTER_METHODDEF + TOKEN_EXIT_METHODDEF {NULL} }; From 0287e78999a3d5bf6fa7e93c4b2920cd475a2496 Mon Sep 17 00:00:00 2001 From: Andrew Svetlov Date: Sun, 9 Feb 2025 13:03:26 +0100 Subject: [PATCH 5/8] More tests --- Lib/test/test_context.py | 27 +++++++++++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/Lib/test/test_context.py b/Lib/test/test_context.py index 33765f45124f9f..874d28aaf71984 100644 --- a/Lib/test/test_context.py +++ b/Lib/test/test_context.py @@ -408,6 +408,33 @@ def fun(): 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 token: + with token: + self.assertEqual(c.get(), 36) + + self.assertEqual(c.get(), 42) + + ctx.run(fun) # HAMT Tests From 4e09df89fec8fc9b86d64879c505455b2f4e1a49 Mon Sep 17 00:00:00 2001 From: Andrew Svetlov Date: Sun, 9 Feb 2025 13:07:36 +0100 Subject: [PATCH 6/8] Fix test --- Lib/test/test_context.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/Lib/test/test_context.py b/Lib/test/test_context.py index 874d28aaf71984..b4f5a2dbfc30c7 100644 --- a/Lib/test/test_context.py +++ b/Lib/test/test_context.py @@ -428,9 +428,13 @@ def test_token_contextmanager_reentrant(self): def fun(): token = c.set(36) - with token: + with self.assertRaisesRegex( + RuntimeError, + " Date: Sun, 9 Feb 2025 13:23:06 +0100 Subject: [PATCH 7/8] More tests --- Lib/test/test_context.py | 53 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 53 insertions(+) diff --git a/Lib/test/test_context.py b/Lib/test/test_context.py index b4f5a2dbfc30c7..f9cdcc3561e9d6 100644 --- a/Lib/test/test_context.py +++ b/Lib/test/test_context.py @@ -440,6 +440,59 @@ def fun(): 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, + " Date: Mon, 10 Feb 2025 09:50:12 +0100 Subject: [PATCH 8/8] Add news --- Doc/whatsnew/3.14.rst | 8 +++++++- .../2025-02-10-09-45-49.gh-issue-129889.PBHXU5.rst | 2 ++ 2 files changed, 9 insertions(+), 1 deletion(-) create mode 100644 Misc/NEWS.d/next/Library/2025-02-10-09-45-49.gh-issue-129889.PBHXU5.rst diff --git a/Doc/whatsnew/3.14.rst b/Doc/whatsnew/3.14.rst index 23b1f5f2e0c67c..fa0f6dae8df5c4 100644 --- a/Doc/whatsnew/3.14.rst +++ b/Doc/whatsnew/3.14.rst @@ -1,4 +1,3 @@ - **************************** What's new in Python 3.14 **************************** @@ -361,6 +360,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/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.