8000 gh-129889: Support context manager protocol by contextvars.Token (#12… · python/cpython@469d2e4 · GitHub
[go: up one dir, main page]
Skip to content

Commit 469d2e4

Browse files
asvetlovasvetlov
authored
gh-129889: Support context manager protocol by contextvars.Token (#129888)
1 parent e1b38ea commit 469d2e4

File tree

6 files changed

+223
-2
lines changed

6 files changed

+223
-2
lines changed

‎Doc/library/contextvars.rst

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -101,6 +101,21 @@ Context Variables
101101
the value of the variable to what it was before the corresponding
102102
*set*.
103103

104+
The token supports :ref:`context manager protocol <context-managers>`
105+
to restore the corresponding context variable value at the exit from
106+
:keyword:`with` block::
107+
108+
var = ContextVar('var', default='default value')
109+
110+
with var.set('new value'):
111+
assert var.get() == 'new value'
112+
113< 10000 /td>+
assert var.get() == 'default value'
114+
115+
.. versionadded:: next
116+
117+
Added support for usage as a context manager.
118+
104119
.. attribute:: Token.var
105120

106121
A read-only property. Points to the :class:`ContextVar` object

‎Doc/whatsnew/3.14.rst

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,3 @@
1-
21
****************************
32
What's new in Python 3.14
43
****************************
@@ -362,6 +361,13 @@ concurrent.futures
362361
supplying a *mp_context* to :class:`concurrent.futures.ProcessPoolExecutor`.
363362
(Contributed by Gregory P. Smith in :gh:`84559`.)
364363

364+
contextvars
365+
-----------
366+
367+
* Support context manager protocol by :class:`contextvars.Token`.
368+
(Contributed by Andrew Svetlov in :gh:`129889`.)
369+
370+
365371
ctypes
366372
------
367373

‎Lib/test/test_context.py

Lines changed: 109 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -383,6 +383,115 @@ def sub(num):
383383
tp.shutdown()
384384
self.assertEqual(results, list(range(10)))
385385

386+
def test_token_contextmanager_with_default(self):
387+
ctx = contextvars.Context()
388+
c = contextvars.ContextVar('c', default=42)
389+
390+
def fun():
391+
with c.set(36):
392+
self.assertEqual(c.get(), 36)
393+
394+
self.assertEqual(c.get(), 42)
395+
396+
ctx.run(fun)
397+
398+
def test_token_contextmanager_without_default(self):
399+
ctx = contextvars.Context()
400+
c = contextvars.ContextVar('c')
401+
402+
def fun():
403+
with c.set(36):
404+
self.assertEqual(c.get(), 36)
405+
406+
with self.assertRaisesRegex(LookupError, "<ContextVar name='c'"):
407+
c.get()
408+
409+
ctx.run(fun)
410+
411+
def test_token_contextmanager_on_exception(self):
412+
ctx = contextvars.Context()
413+
c = contextvars.ContextVar('c', default=42)
414+
415+
def fun():
416+
with c.set(36):
417+
self.assertEqual(c.get(), 36)
418+
raise ValueError("custom exception")
419+
420+
self.assertEqual(c.get(), 42)
421+
422+
with self.assertRaisesRegex(ValueError, "custom exception"):
423+
ctx.run(fun)
424+
425+
def test_token_contextmanager_reentrant(self):
426+
ctx = contextvars.Context()
427+
c = contextvars.ContextVar('c', default=42)
428+
429+
def fun():
430+
token = c.set(36)
431+
with self.assertRaisesRegex(
432+
RuntimeError,
433+
"<Token .+ has already been used once"
434+
):
435+
with token:
436+
with token:
437+
self.assertEqual(c.get(), 36)
438+
439+
self.assertEqual(c.get(), 42)
440+
441+
ctx.run(fun)
442+
443+
def test_token_contextmanager_multiple_c_set(self):
444+
ctx = contextvars.Context()
445+
c = contextvars.ContextVar('c', default=42)
446+
447+
def fun():
448+
with c.set(36):
449+
self.assertEqual(c.get(), 36)
450+
c.set(24)
451+
self.assertEqual(c.get(), 24)
452+
c.set(12)
453+
self.assertEqual(c.get(), 12)
454+
455+
self.assertEqual(c.get(), 42)
456+
457+
ctx.run(fun)
458+
459+
def test_token_contextmanager_with_explicit_reset_the_same_token(self):
460+
ctx = contextvars.Context()
461+
c = contextvars.ContextVar('c', default=42)
462+
463+
def fun():
464+
with self.assertRaisesRegex(
465+
RuntimeError,
466+
"<Token .+ has already been used once"
467+
):
468+
with c.set(36) as token:
469+
self.assertEqual(c.get(), 36)
470+
c.reset(token)
471+
472+
self.assertEqual(c.get(), 42)
473+
474+
self.assertEqual(c.get(), 42)
475+
476+
ctx.run(fun)
477+
478+
def test_token_contextmanager_with_explicit_reset_another_token(self):
479+
ctx = contextvars.Context()
480+
c = contextvars.ContextVar('c', default=42)
481+
482+
def fun():
483+
with c.set(36):
484+
self.assertEqual(c.get(), 36)
485+
486+
token = c.set(24)
487+
self.assertEqual(c.get(), 24)
488+
c.reset(token)
489+
self.assertEqual(c.get(), 36)
490+
491+
self.assertEqual(c.get(), 42)
492+
493+
ctx.run(fun)
494+
386495

387496
# HAMT Tests
388497

‎Misc/NEWS.d/next/Library/2025-02-10-09-45-49.gh-issue-129889.PBHXU5.rst

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
Support context manager protocol by :class:`contextvars.Token`. Patch by
2+
Andrew Svetlov.

‎Python/clinic/context.c.h

Lines changed: 52 additions & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎Python/context.c

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1231,9 +1231,47 @@ static PyGetSetDef PyContextTokenType_getsetlist[] = {
12311231
{NULL}
12321232
};
12331233

1234+
/*[clinic input]
1235+
_contextvars.Token.__enter__ as token_enter
1236+
1237+
Enter into Token context manager.
1238+
[clinic start generated code]*/
1239+
1240+
static PyObject *
1241+
token_enter_impl(PyContextToken *self)
1242+
/*[clinic end generated code: output=9af4d2054e93fb75 input=41a3d6c4195fd47a]*/
1243+
{
1244+
return Py_NewRef(self);
1245+
}
1246+
1247+
/*[clinic input]
1248+
_contextvars.Token.__exit__ as token_exit
1249+
1250+
type: object
1251+
val: object
1252+
tb: object
1253+
/
1254+
1255+
Exit from Token context manager, restore the linked ContextVar.
1256+
[clinic start generated code]*/
1257+
1258+
static PyObject *
1259+
token_exit_impl(PyContextToken *self, PyObject *type, PyObject *val,
1260+
PyObject *tb)
1261+
/*[clinic end generated code: output=3e6a1c95d3da703a input=7f117445f0ccd92e]*/
1262+
{
1263+
int ret = PyContextVar_Reset((PyObject *)self->tok_var, (PyObject *)self);
1264+
if (ret < 0) {
1265+
return NULL;
1266+
}
1267+
Py_RETURN_NONE;
1268+
}
1269+
12341270
static PyMethodDef PyContextTokenType_methods[] = {
12351271
{"__class_getitem__", Py_GenericAlias,
12361272
METH_O|METH_CLASS, PyDoc_STR("See PEP 585")},
1273+
TOKEN_ENTER_METHODDEF
1274+
TOKEN_EXIT_METHODDEF
12371275
{NULL}
12381276
};
12391277

0 commit comments

Comments
 (0)
0