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

Commit f5ef259

Browse files
rhansenrhansen
committed
gh-99633: Add context manager support to contextvars.Context 
1 parent ee21110 commit f5ef259

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.12
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
@@ -360,6 +361,74 @@ def sub(num):
360361
tp.shutdown()
361362
self.assertEqual(results, list(range(10)))
362363

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

364433
# HAMT Tests
365434

‎Misc/ACKS

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -694,6 +694,7 @@ Manus Hand
694694
Andreas Hangauer
695695
Milton L. Hankins
696696
Carl Bordum Hansen
697+
Richard Hansen
697698
Stephen Hansen
698699
Barry Hantman
699700
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");
@@ -558,6 +557,69 @@ context_tp_contains(PyContext *self, PyObject *key)
558557
}
559558

560559

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

678740

679741
static PyMethodDef PyContext_methods[] = {
742+
_CONTEXTVARS_CONTEXT___ENTER___METHODDEF
743+
_CONTEXTVARS_CONTEXT___EXIT___METHODDEF
680744
_CONTEXTVARS_CONTEXT_GET_METHODDEF
681745
_CONTEXTVARS_CONTEXT_ITEMS_METHODDEF
682746
_CONTEXTVARS_CONTEXT_KEYS_METHODDEF

0 commit comments

Comments
 (0)
0