10000 gh-99633: Add context manager support to `contextvars.Context` · python/cpython@442d790 · GitHub
[go: up one dir, main page]
Skip to content

Commit 442d790

Browse files
rhansenrhansen
committed
gh-99633: Add context manager support to contextvars.Context 
1 parent 3ec719f commit 442d790

File tree

6 files changed

+271
-15
lines changed

6 files changed

+271
-15
lines changed

‎Doc/library/contextvars.rst

Lines changed: 67 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -144,21 +144,79 @@ Manual Context Management
144144
To get a copy of the current context use the
145145
:func:`~contextvars.copy_context` function.
146146

147-
Every thread will have a different top-level :class:`~contextvars.Context`
148-
object. This means that a :class:`ContextVar` object behaves in a similar
149-
fashion to :func:`threading.local()` when values are assigned in different
150-
threads.
147+
Each thread has its own effective stack of :class:`~contextvars.Context`
148+
objects. The *current Context* is the Context object at the top of the
149+
current thread's stack. All Context objects in the stacks are considered to
150+
be *entered*. *Entering* a Context, either by calling the
151+
:meth:`Context.run` method or using the Context object as a :term:`context
152+
manager`, pushes the Context onto the top of the current thread's stack,
153+
making it the current Context. *Exiting* from the current Context, either by
154+
returning from the callback passed to :meth:`Context.run` or by exiting the
155+
:keyword:`with` statement suite, pops the Context off of the top of the
156+
stack, restoring the current Context to what it was before.
157+
158+
Because each thread has its own Context stack, :class:`ContextVar` objects
159+
behave in a similar fashion to :func:`threading.local()` when values are
160+
assigned in different threads.
161+
162+
Attempting to do either of the following will raise a :exc:`RuntimeError`:
163+
164+
* Entering an already entered Context. (This includes Contexts entered in
165+
other threads.)
166+
* Exiting from a Context that is not the current Context.
167+
168+
After exiting a Context, it can later be re-entered (from any thread).
169+
170+
Any changes to :class:`ContextVar` values via the :meth:`ContextVar.set`
171+
method are recorded in the current Context. The :meth:`ContextVar.get`
172+
method returns the value associated with the current Context. Thus, exiting
173+
a Context effectively reverts any changes made to context variables while the
174+
Context was entered. (If desired, the values can be restored by re-entering
175+
the Context.)
151176

152177
Context implements the :class:`collections.abc.Mapping` interface.
153178

179+
.. versionadded:: 3.14
180+
A Context object can be used as a :term:`context manager`. The
181+
:meth:`Context.__enter__` and :meth:`Context.__exit__` methods
182+
(automatically called by the :keyword:`with` statement) enters and exits
183+
the Context, respectively. The value returned from
184+
:meth:`Context.__enter__`, and thus bound to the identifier given in the
185+
:keyword:`with` statement's :keyword:`!as` clause if present, is the
186+
Context object itself.
187+
188+
Example:
189+
190+
.. testcode::
191+
192+
import contextvars
193+
194+
var = contextvars.ContextVar("var")
195+
var.set("initial")
196+
197+
# Copy the current Context and enter it.
198+
with contextvars.copy_context() as ctx:
199+
var.set("updated")
200+
assert var in ctx
201+
assert ctx[var] == "updated"
202+
assert var.get() == "updated"
203+
204+
# Exited ctx, so the value of var should have reverted.
205+
assert var.get() == "initial"
206+
# But the updated value is still recorded in ctx.
207+
assert ctx[var] == "updated"
208+
209+
# Re-entering ctx should restore the updated value of var.
210+
with ctx:
211+
assert var.get() == "updated"
212+
154213
.. method:: run(callable, *args, **kwargs)
155214

156-
Execute ``callable(*args, **kwargs)`` code in the context object
157-
the *run* method is called on. Return the result of the execution
158-
or propagate an exception if one occurred.
215+
Enters the Context, executes ``callable(*args, **kwargs)``, then exits the
216+
Context. Returns *callable*'s return value, or propagates an exception if
217+
one occurred.
159218

160-
Any changes to any context variables that *callable* makes will
161-
be contained in the context object::
219+
Example::
162220

163221
var = ContextVar('var')
164222
var.set('spam')
@@ -186,10 +244,6 @@ Manual Context Management
186244
# However, outside of 'ctx', 'var' is still set to 'spam':
187245
# var.get() == 'spam'
188246

189-
The method raises a :exc:`RuntimeError` when called on the same
190-
context object from more than one OS thread, or when called
191-
recursively.
192-
193247
.. method:: copy()
194248

195249
Return a shallow copy of the context object.

‎Lib/test/test_context.py

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
import functools
44
import gc
55
import random
6+
import threading
67
import time
78
import unittest
89
import weakref
@@ -361,6 +362,74 @@ def sub(num):
361362
tp.shutdown()
362363
self.assertEqual(results, list(range(10)))
363364

365+
@isolated_context
366+
def test_context_manager_works(self):
367+
cvar = contextvars.ContextVar('cvar', default='initial')
368+
self.assertEqual(cvar.get(), 'initial')
369+
with contextvars.copy_context():
370+
self.assertEqual(cvar.get(), 'initial')
371+
cvar.set('updated')
372+
self.assertEqual(cvar.get(), 'updated')
373+
self.assertEqual(cvar.get(), 'initial')
374+
375+
def test_context_manager_as_binding(self):
376+
ctx = contextvars.copy_context()
377+
with ctx as ctx_as_binding:
378+
self.assertIs(ctx_as_binding, ctx)
379+
380+
@isolated_context
381+
def test_context_manager_enter_again_after_exit(self):
382+
cvar = contextvars.ContextVar('cvar', default='initial')
383+
self.assertEqual(cvar.get(), 'initial')
384+
with contextvars.copy_context() as ctx:
385+
cvar.set('updated')
386+
self.assertEqual(cvar.get(), 'updated')
387+
self.assertEqual(cvar.get(), 'initial')
388+
with ctx:
389+
self.assertEqual(cvar.get(), 'updated')
390+
self.assertEqual(cvar.get(), 'initial')
391+
392+
@threading_helper.requires_working_threading()
393+
def test_context_manager_rejects_exit_from_different_thread(self):
394+
ctx = contextvars.copy_context()
395+
thread = threading.Thread(target=ctx.__enter__)
396+
thread.start()
397+
thread.join()
398+
with self.assertRaises(RuntimeError):
399+
ctx.__exit__(None, None, None)
400+
401+
def test_context_manager_rejects_recursive_enter_mgr_then_mgr(self):
402+
with contextvars.copy_context() as ctx:
403+
with self.assertRaises(RuntimeError):
404+
with ctx:
405+
pass
406+
407+
def test_context_manager_rejects_recursive_enter_mgr_then_run(self):
408+
with contextvars.copy_context() as ctx:
409+
with self.assertRaises(RuntimeError):
410+
ctx.run(lambda: None)
411+
412+
def test_context_manager_rejects_recursive_enter_run_then_mgr(self):
413+
ctx = contextvars.copy_context()
414+
415+
def fn():
416+
with self.assertRaises(RuntimeError):
417+
with ctx:
418+
pass
419+
420+
ctx.run(fn)
421+
422+
def test_context_manager_rejects_noncurrent_exit(self):
423+
with contextvars.copy_context() as ctx:
424+
with contextvars.copy_context():
425+
with self.assertRaises(RuntimeError):
426+
ctx.__exit__(None, None, None)
427+
428+
def test_context_manager_rejects_nonentered_exit(self):
429+
ctx = contextvars.copy_context()
430+
with self.assertRaises(RuntimeError):
431+
ctx.__exit__(None, None, None)
432+
364433

365434
# HAMT Tests
366435

‎Misc/ACKS

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -715,6 +715,7 @@ Michael Handler
715715
Andreas Hangauer
716716
Milton L. Hankins
717717
Carl Bordum Hansen
718+
Richard Hansen
718719
Stephen Hansen
719720
Barry Hantman
720721
Lynda Hardman

‎Misc/NEWS.d/next/Library/2022-11-21-01-24-46.gh-issue-99633.vhrNRe.rst

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Add :term:`context manager` methods to :class:`contextvars.Context`.

‎Python/clinic/context.c.h

Lines changed: 68 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: 65 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -153,7 +153,6 @@ _PyContext_Exit(PyThreadState *ts, PyObject *octx)
153153
}
154154

155155
if (ts->context != (PyObject *)ctx) {
156-
/* Can only happen if someone misuses the C API */
157156
PyErr_SetString(PyExc_RuntimeError,
158157
"cannot exit context: thread state references "
159158
"a different context object");
@@ -550,6 +549,69 @@ context_tp_contains(PyContext *self, PyObject *key)
550549
}
551550

552551

552+
/*[clinic input]
553+
_contextvars.Context.__enter__
554+
555+
Context manager enter.
556+
557+
Automatically called by the 'with' statement. Using the Context object as a
558+
context manager is an alternative to calling the Context.run() method. Example
559+
usage:
560+
561+
var = contextvars.ContextVar('var')
562+
var.set('initial')
563+
564+
with contextvars.copy_context():
565+
# The current Context is a new copy of the previous Context. Updating a
566+
# context variable inside this 'with' statement only affects the new
567+
# copy.
568+
var.set('updated')
569+
do_something_interesting()
570+
571+
# Now that the 'with' statement is done executing, the value of the 'var'
572+
# context variable has reverted back to its value before the 'with'
573+
# statement.
574+
assert var.get() == 'initial'
575+
[clinic start generated code]*/
576+
577+
static PyObject *
578+
_contextvars_Context___enter___impl(PyContext *self)
579+
/*[clinic end generated code: output=7374aea8983b777a input=22868e8274d3cd32]*/
580+
{
581+
PyThreadState *ts = _PyThreadState_GET();
582+
if (_PyContext_Enter(ts, (PyObject *)self)) {
583+
return NULL;
584+
}
585+
return Py_NewRef(self);
586+
}
587+
588+
589+
/*[clinic input]
590+
_contextvars.Context.__exit__
591+
exc_type: object
592+
exc_val: object
593+
exc_tb: object
594+
/
595+
596+
Context manager exit.
597+
598+
Automatically called at the conclusion of a 'with' statement when the Context is
599+
used as a context manager. See the Context.__enter__() method for more details.
600+
[clinic start generated code]*/
601+
602+
static PyObject *
603+
_contextvars_Context___exit___impl(PyContext *self, PyObject *exc_type,
604+
PyObject *exc_val, PyObject *exc_tb)
605+
/*[clinic end generated code: output=4608fa9151f968f1 input=ff70cbbf6a112b1d]*/
606+
{
607+
PyThreadState *ts = _PyThreadState_GET();
608+
if (_PyContext_Exit(ts, (PyObject *)self)) {
609+
return NULL;
610+
}
611+
Py_RETURN_NONE;
612+
}
613+
614+
553615
/*[clinic input]
554616
_contextvars.Context.get
555617
key: object
@@ -670,6 +732,8 @@ context_run(PyContext *self, PyObject *const *args,
670732

671733

672734
static PyMethodDef PyContext_methods[] = {
735+
_CONTEXTVARS_CONTEXT___ENTER___METHODDEF
736+
_CONTEXTVARS_CONTEXT___EXIT___METHODDEF
673737
_CONTEXTVARS_CONTEXT_GET_METHODDEF
674738
_CONTEXTVARS_CONTEXT_ITEMS_METHODDEF
675739
_CONTEXTVARS_CONTEXT_KEYS_METHODDEF

0 commit comments

Comments
 (0)
0