Skip to content

Commit 0678535

Browse files
committed
gh-99633: Add context manager support to contextvars.Context
1 parent e393c79 commit 0678535

File tree

7 files changed

+304
-16
lines changed

7 files changed

+304
-16
lines changed

Doc/library/contextvars.rst

Lines changed: 79 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -144,21 +144,90 @@ 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:`!Context` objects. The
148+
*current context* is the :class:`!Context` object at the top of the current
149+
thread's stack. All :class:`!Context` objects in the stacks are considered
150+
to be *entered*.
151+
152+
*Entering* a context, either by calling the :meth:`~Context.run` method or
153+
using the context as a :term:`context manager`, pushes the context onto the
154+
top of the current thread's stack, making it the current context.
155+
156+
*Exiting* from the current context, either by returning from the callback
157+
passed to :meth:`~Context.run` or by exiting the :keyword:`with` statement
158+
suite, pops the context off of the top of the stack, restoring the current
159+
context to what it was before.
160+
161+
Since each thread has its own context stack, :class:`ContextVar` objects
162+
behave in a similar fashion to :func:`threading.local` when values are
163+
assigned in different threads.
164+
165+
Attempting to do either of the following raises a :exc:`RuntimeError`:
166+
167+
* Entering an already entered context, including contexts entered in
168+
other threads.
169+
* Exiting from a context that is not the current context.
170+
171+
After exiting a context, it can later be re-entered (from any thread).
172+
173+
Any changes to :class:`ContextVar` values via the :meth:`ContextVar.set`
174+
method are recorded in the current context. The :meth:`ContextVar.get`
175+
method returns the value associated with the current context. Exiting a
176+
context effectively reverts any changes made to context variables while the
177+
context was entered (if needed, the values can be restored by re-entering the
178+
context).
151179

152180
Context implements the :class:`collections.abc.Mapping` interface.
153181

154-
.. method:: run(callable, *args, **kwargs)
182+
.. versionadded:: 3.14
183+
Added support for the :term:`context management protocol`.
184+
185+
When used as a :term:`context manager`, the value bound to the identifier
186+
given in the :keyword:`with` statement's :keyword:`!as` clause (if present)
187+
is the :class:`!Context` object itself.
188+
189+
Example:
155190

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.
191+
.. testcode::
159192

160-
Any changes to any context variables that *callable* makes will
161-
be contained in the context object.
193+
import contextvars
194+
195+
var = contextvars.ContextVar("var")
196+
var.set("initial")
197+
print(var.get()) # 'initial'
198+
199+
# Copy the current Context and enter the copy.
200+
with contextvars.copy_context() as ctx:
201+
var.set("updated")
202+
print(var in ctx) # 'True'
203+
print(ctx[var]) # 'updated'
204+
print(var.get()) # 'updated'
205+
206+
# Exited ctx, so the observed value of var has reverted.
207+
print(var.get()) # 'initial'
208+
# But the updated value is still recorded in ctx.
209+
print(ctx[var]) # 'updated'
210+
211+
# Re-entering ctx restores the updated value of var.
212+
with ctx:
213+
print(var.get()) # 'updated'
214+
215+
.. testoutput::
216+
:hide:
217+
218+
initial
219+
True
220+
updated
221+
updated
222+
initial
223+
updated
224+
updated
225+
226+
.. method:: run(callable, *args, **kwargs)
227+
228+
Enters the Context, executes ``callable(*args, **kwargs)``, then exits the
229+
Context. Returns *callable*'s return value, or propagates an exception if
230+
one occurred.
162231

163232
Example:
164233

@@ -206,10 +275,6 @@ Manual Context Management
206275
ham
207276
spam
208277

209-
The method raises a :exc:`RuntimeError` when called on the same
210-
context object from more than one OS thread, or when called
211-
recursively.
212-
213278
.. method:: copy()
214279

215280
Return a shallow copy of the context object.

Doc/whatsnew/3.14.rst

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -129,6 +129,13 @@ ast
129129
(Contributed by Irit Katriel in :gh:`123958`.)
130130

131131

132+
contextvars
133+
-----------
134+
135+
* Added support for the :term:`context management protocol` to
136+
:class:`contextvars.Context`. (Contributed by Richard Hansen in :gh:`99634`.)
137+
138+
132139
ctypes
133140
------
134141

Lib/test/test_context.py

Lines changed: 83 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,88 @@ def sub(num):
361362
tp.shutdown()
362363
self.assertEqual(results, list(range(10)))
363364

365+
@isolated_context
366+
def test_context_manager(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_nested(self):
382+
cvar = contextvars.ContextVar('cvar', default='default')
383+
with contextvars.copy_context() as outer_ctx:
384+
cvar.set('outer')
385+
with contextvars.copy_context() as inner_ctx:
386+
self.assertIsNot(outer_ctx, inner_ctx)
387+
self.assertEqual(cvar.get(), 'outer')
388+
cvar.set('inner')
389+
self.assertEqual(outer_ctx[cvar], 'outer')
390+
self.assertEqual(cvar.get(), 'inner')
391+
self.assertEqual(cvar.get(), 'outer')
392+
self.assertEqual(cvar.get(), 'default')
393+
394+
@isolated_context
395+
def test_context_manager_enter_again_after_exit(self):
396+
cvar = contextvars.ContextVar('cvar', default='initial')
397+
self.assertEqual(cvar.get(), 'initial')
398+
with contextvars.copy_context() as ctx:
399+
cvar.set('updated')
400+
self.assertEqual(cvar.get(), 'updated')
401+
self.assertEqual(cvar.get(), 'initial')
402+
with ctx:
403+
self.assertEqual(cvar.get(), 'updated')
404+
self.assertEqual(cvar.get(), 'initial')
405+
406+
@threading_helper.requires_working_threading()
407+
def test_context_manager_rejects_exit_from_different_thread(self):
408+
ctx = contextvars.copy_context()
409+
thread = threading.Thread(target=ctx.__enter__)
410+
thread.start()
411+
thread.join()
412+
with self.assertRaises(RuntimeError):
413+
ctx.__exit__(None, None, None)
414+
415+
def test_context_manager_is_not_reentrant(self):
416+
with self.subTest('context manager then context manager'):
417+
with contextvars.copy_context() as ctx:
418+
with self.assertRaises(RuntimeError):
419+
with ctx:
420+
pass
421+
with self.subTest('context manager then run method'):
422+
with contextvars.copy_context() as ctx:
423+
with self.assertRaises(RuntimeError):
424+
ctx.run(lambda: None)
425+
with self.subTest('run method then context manager'):
426+
ctx = contextvars.copy_context()
427+
428+
def fn():
429+
with self.assertRaises(RuntimeError):
430+
with ctx:
431+
pass
432+
433+
ctx.run(fn)
434+
435+
def test_context_manager_rejects_noncurrent_exit(self):
436+
with contextvars.copy_context() as outer_ctx:
437+
with contextvars.copy_context() as inner_ctx:
438+
self.assertIsNot(outer_ctx, inner_ctx)
439+
with self.assertRaises(RuntimeError):
440+
outer_ctx.__exit__(None, None, None)
441+
442+
def test_context_manager_rejects_nonentered_exit(self):
443+
ctx = contextvars.copy_context()
444+
with self.assertRaises(RuntimeError):
445+
ctx.__exit__(None, None, None)
446+
364447

365448
# HAMT Tests
366449

Misc/ACKS

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -716,6 +716,7 @@ Michael Handler
716716
Andreas Hangauer
717717
Milton L. Hankins
718718
Carl Bordum Hansen
719+
Richard Hansen
719720
Stephen Hansen
720721
Barry Hantman
721722
Lynda Hardman
Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
Added support for the :term:`context management protocol` to
2+
:class:`contextvars.Context`. Patch by Richard Hansen.

Python/clinic/context.c.h

Lines changed: 69 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: 63 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -144,7 +144,6 @@ _PyContext_Exit(PyThreadState *ts, PyObject *octx)
144144
}
145145

146146
if (ts->context != (PyObject *)ctx) {
147-
/* Can only happen if someone misuses the C API */
148147
PyErr_SetString(PyExc_RuntimeError,
149148
"cannot exit context: thread state references "
150149
"a different context object");
@@ -521,6 +520,67 @@ context_tp_contains(PyContext *self, PyObject *key)
521520
}
522521

523522

523+
/*[clinic input]
524+
_contextvars.Context.__enter__
525+
526+
Context manager enter.
527+
528+
Automatically called by the 'with' statement. Using the Context object as a
529+
context manager is an alternative to calling the Context.run() method.
530+
531+
Example:
532+
533+
var = contextvars.ContextVar('var')
534+
var.set('initial')
535+
536+
with contextvars.copy_context():
537+
# Changes to context variables will be rolled back upon exiting the
538+
# `with` statement.
539+
var.set('updated')
540+
print(var.get()) # 'updated'
541+
542+
# The context variable value has been rolled back.
543+
print(var.get()) # 'initial'
544+
[clinic start generated code]*/
545+
546+
static PyObject *
547+
_contextvars_Context___enter___impl(PyContext *self)
548+
/*[clinic end generated code: output=7374aea8983b777a input=fffe71e56ca17ee4]*/
549+
{
550+
// The new ref added here is for the `with` statement's `as` binding. It is
551+
// decremented when the variable goes out of scope, which can be before or
552+
// after `PyContext_Exit` is called. (The binding can go out of scope
553+
// immediately -- before the `with` suite even runs -- if there is no `as`
554+
// clause. Or it can go out of scope long after the `with` suite completes
555+
// because `with` does not have its own scope.) Because of this timing, two
556+
// references are needed: the one added in `PyContext_Enter` and the one
557+
// added here.
558+
return PyContext_Enter((PyObject *)self) < 0 ? NULL : Py_NewRef(self);
559+
}
560+
561+
562+
/*[clinic input]
563+
_contextvars.Context.__exit__
564+
exc_type: object
565+
exc_val: object
566+
exc_tb: object
567+
/
568+
569+
Context manager exit.
570+
571+
Automatically called at the conclusion of a 'with' statement when the Context is
572+
used as a context manager. See the Context.__enter__() method for more details.
573+
[clinic start generated code]*/
574+
575+
static PyObject *
576+
_contextvars_Context___exit___impl(PyContext *self, PyObject *exc_type,
577+
PyObject *exc_val, PyObject *exc_tb)
578+
/*[clinic end generated code: output=4608fa9151f968f1 input=ff70cbbf6a112b1d]*/
579+
{
580+
return PyContext_Exit((PyObject *)self) < 0 ? NULL : Py_None;
581+
}
582+
583+
524584
/*[clinic input]
525585
_contextvars.Context.get
526586
key: object
@@ -641,6 +701,8 @@ context_run(PyContext *self, PyObject *const *args,
641701

642702

643703
static PyMethodDef PyContext_methods[] = {
704+
_CONTEXTVARS_CONTEXT___ENTER___METHODDEF
705+
_CONTEXTVARS_CONTEXT___EXIT___METHODDEF
644706
_CONTEXTVARS_CONTEXT_GET_METHODDEF
645707
_CONTEXTVARS_CONTEXT_ITEMS_METHODDEF
646708
_CONTEXTVARS_CONTEXT_KEYS_METHODDEF

0 commit comments

Comments
 (0)