diff --git a/Doc/c-api/contextvars.rst b/Doc/c-api/contextvars.rst index b7c6550ff34aac..7489ea7bf28810 100644 --- a/Doc/c-api/contextvars.rst +++ b/Doc/c-api/contextvars.rst @@ -155,7 +155,37 @@ Context variable functions: Create a new ``ContextVar`` object. The *name* parameter is used for introspection and debug purposes. The *def* parameter specifies a default value for the context variable, or ``NULL`` for no default. - If an error has occurred, this function returns ``NULL``. + Automatic inheritance of the variable's bindings by + :class:`threading.Thread` follows + :data:`sys.flags.thread_inherit_context`. If an error has occurred, this + function returns ``NULL``. + +.. c:function:: PyObject *PyContextVar_NewWithFlags(const char *name, PyObject *def, int flags) + + Create a new ``ContextVar`` object with the specified *flags*. The *name*, + *def*, and return value are the same as for :c:func:`PyContextVar_New`. + *flags* must be one of the following values: + + .. c:macro:: Py_CONTEXTVAR_INHERIT_THREAD_DEFAULT + + Follow :data:`sys.flags.thread_inherit_context` when deciding whether a + new :class:`threading.Thread` inherits the variable's binding. This is + equivalent to :c:func:`PyContextVar_New`. + + .. c:macro:: Py_CONTEXTVAR_INHERIT_THREAD_NEVER + + Do not automatically inherit the variable's binding, regardless of + :data:`sys.flags.thread_inherit_context`. + + .. c:macro:: Py_CONTEXTVAR_INHERIT_THREAD_ALWAYS + + Automatically inherit the variable's binding, regardless of + :data:`sys.flags.thread_inherit_context`. + + Passing any other value causes the function to return ``NULL`` with + :exc:`ValueError` set. + + .. versionadded:: 3.16 .. c:function:: int PyContextVar_Get(PyObject *var, PyObject *default_value, PyObject **value) diff --git a/Doc/data/refcounts.dat b/Doc/data/refcounts.dat index 60c02aabeb89c5..7b8d5fc2335f21 100644 --- a/Doc/data/refcounts.dat +++ b/Doc/data/refcounts.dat @@ -392,6 +392,11 @@ PyContextVar_New:PyObject*::+1: PyContextVar_New:const char*:name:: PyContextVar_New:PyObject*:def:+1: +PyContextVar_NewWithFlags:PyObject*::+1: +PyContextVar_NewWithFlags:const char*:name:: +PyContextVar_NewWithFlags:PyObject*:def:+1: +PyContextVar_NewWithFlags:int:flags:: + PyContextVar_Set:PyObject*::+1: PyContextVar_Set:PyObject*:var:0: PyContextVar_Set:PyObject*:value:+1: diff --git a/Doc/howto/free-threading-python.rst b/Doc/howto/free-threading-python.rst index 53bea1db191d76..91a7a4359cc253 100644 --- a/Doc/howto/free-threading-python.rst +++ b/Doc/howto/free-threading-python.rst @@ -152,8 +152,9 @@ is set to true by default which causes threads created with :class:`threading.Thread` to start with a copy of the :class:`~contextvars.Context()` of the caller of :meth:`~threading.Thread.start`. In the default GIL-enabled build, the flag -defaults to false so threads start with an -empty :class:`~contextvars.Context()`. +defaults to false, so threads do not inherit context variable bindings by +default. Individual variables can override either setting with the +:class:`~contextvars.ContextVar` constructor's *thread_inheritable* parameter. Warning filters diff --git a/Doc/library/contextvars.rst b/Doc/library/contextvars.rst index b0cc0be8e911bf..4280312c797dc1 100644 --- a/Doc/library/contextvars.rst +++ b/Doc/library/contextvars.rst @@ -24,7 +24,7 @@ See also :pep:`567` for additional details. Context Variables ----------------- -.. class:: ContextVar(name, [*, default]) +.. class:: ContextVar(name, [*, default, thread_inheritable]) This class is used to declare a new Context Variable, e.g.:: @@ -37,6 +37,38 @@ Context Variables :meth:`ContextVar.get` when no value for the variable is found in the current context. + The optional keyword-only *thread_inheritable* parameter controls whether + the variable's current binding is automatically inherited by a + :class:`threading.Thread` when no explicit :class:`Context` is supplied. + If ``None`` (the default), :data:`sys.flags.thread_inherit_context` + determines whether the binding is inherited. If true, the binding is + inherited regardless of that flag; if false, it is not inherited regardless + of the flag. + + This parameter only affects implicit context inheritance by + :meth:`threading.Thread.start`. The bindings are captured from the caller + of :meth:`~threading.Thread.start`, not the creator of the + :class:`~threading.Thread` object. It does not affect threads started + directly through :mod:`_thread` or by extension modules. + + An explicitly supplied thread context is not filtered, and neither are + contexts copied for other purposes, including asynchronous tasks and + :func:`asyncio.to_thread`. An inherited binding is an ordinary binding in + the new thread's context: it is visible to :func:`copy_context`, may be + changed independently, and may in turn be inherited by threads started from + that thread. Only the binding is copied; the bound object itself is not + copied. A mutable bound object can therefore be accessed concurrently by + the starting and new threads. + + For a thread pool, inheritance occurs when each worker thread starts, not + each time work is submitted. Consequently, + :class:`~concurrent.futures.ThreadPoolExecutor` does not propagate the + submitter's current bindings to an existing worker merely because a variable + is thread-inheritable. + + .. versionchanged:: 3.16 + Added the *thread_inheritable* parameter. + **Important:** Context Variables should be created at the top module level and never in closures. :class:`Context` objects hold strong references to context variables which prevents context variables @@ -51,6 +83,13 @@ Context Variables .. versionadded:: 3.7.1 + .. attribute:: ContextVar.thread_inheritable + + The value of the *thread_inheritable* constructor parameter: ``None``, + ``True``, or ``False``. This is a read-only property. + + .. versionadded:: 3.16 + .. method:: get([default]) Return a value for the context variable for the current context. diff --git a/Doc/library/threading.rst b/Doc/library/threading.rst index 5d9a7b6314b166..627e412673e51d 100644 --- a/Doc/library/threading.rst +++ b/Doc/library/threading.rst @@ -539,11 +539,14 @@ since it is impossible to detect the termination of alien threads. the thread. The default value is ``None`` which indicates that the :data:`sys.flags.thread_inherit_context` flag controls the behaviour. If the flag is true, threads will start with a copy of the context of the - caller of :meth:`~Thread.start`. If false, they will start with an empty - context. To explicitly start with an empty context, pass a new instance of - :class:`~contextvars.Context()`. To explicitly start with a copy of the - current context, pass the value from :func:`~contextvars.copy_context`. The - flag defaults true on free-threaded builds and false otherwise. + caller of :meth:`~Thread.start`; if false, they will not inherit bindings by + default. The :class:`~contextvars.ContextVar` constructor's + *thread_inheritable* parameter can override the flag for an individual + variable. To explicitly start with an empty context, pass a new instance + of :class:`~contextvars.Context()`. To explicitly start with a copy of the + current context, pass the value from :func:`~contextvars.copy_context`. + Explicit contexts are not filtered by per-variable settings. The flag + defaults true on free-threaded builds and false otherwise. If the subclass overrides the constructor, it must make sure to invoke the base class constructor (``Thread.__init__()``) before doing anything else to diff --git a/Doc/using/cmdline.rst b/Doc/using/cmdline.rst index 677fbbae3f4219..11f8e139c0bec4 100644 --- a/Doc/using/cmdline.rst +++ b/Doc/using/cmdline.rst @@ -674,12 +674,13 @@ Miscellaneous options .. versionadded:: 3.13 - * :samp:`-X thread_inherit_context={0,1}` causes :class:`~threading.Thread` - to, by default, use a copy of context of the caller of - ``Thread.start()`` when starting. Otherwise, threads will start - with an empty context. If unset, the value of this option defaults - to ``1`` on free-threaded builds and to ``0`` otherwise. See also - :envvar:`PYTHON_THREAD_INHERIT_CONTEXT`. + * :samp:`-X thread_inherit_context={0,1}` controls whether + :class:`~threading.Thread` inherits context variable bindings from the + caller of ``Thread.start()`` by default. Individual context variables can + override this setting with the :class:`~contextvars.ContextVar` + constructor's *thread_inheritable* parameter. If unset, the value of this + option defaults to ``1`` on free-threaded builds and to ``0`` otherwise. + See also :envvar:`PYTHON_THREAD_INHERIT_CONTEXT`. .. versionadded:: 3.14 @@ -1366,11 +1367,13 @@ conflict. .. envvar:: PYTHON_THREAD_INHERIT_CONTEXT - If this variable is set to ``1`` then :class:`~threading.Thread` will, - by default, use a copy of context of the caller of ``Thread.start()`` - when starting. Otherwise, new threads will start with an empty context. - If unset, this variable defaults to ``1`` on free-threaded builds and to - ``0`` otherwise. See also :option:`-X thread_inherit_context<-X>`. + This variable controls whether :class:`~threading.Thread` inherits context + variable bindings from the caller of ``Thread.start()`` by default. + Individual context variables can override it with the + :class:`~contextvars.ContextVar` constructor's *thread_inheritable* + parameter. If unset, this variable defaults to ``1`` on free-threaded + builds and to ``0`` otherwise. See also + :option:`-X thread_inherit_context<-X>`. .. versionadded:: 3.14 diff --git a/Include/cpython/context.h b/Include/cpython/context.h index 3a7a4b459c09ad..471d4ad6c2b8a7 100644 --- a/Include/cpython/context.h +++ b/Include/cpython/context.h @@ -68,6 +68,19 @@ PyAPI_FUNC(int) PyContext_ClearWatcher(int watcher_id); PyAPI_FUNC(PyObject *) PyContextVar_New( const char *name, PyObject *default_value); +/* Flags for PyContextVar_NewWithFlags(). */ +#define Py_CONTEXTVAR_INHERIT_THREAD_DEFAULT 0 +#define Py_CONTEXTVAR_INHERIT_THREAD_NEVER 1 +#define Py_CONTEXTVAR_INHERIT_THREAD_ALWAYS 2 + +/* Create a new context variable with the specified flags. + + default_value can be NULL. flags must be one of the + Py_CONTEXTVAR_INHERIT_THREAD_* values. +*/ +PyAPI_FUNC(PyObject *) PyContextVar_NewWithFlags( + const char *name, PyObject *default_value, int flags); + /* Get a value for the variable. diff --git a/Include/internal/pycore_context.h b/Include/internal/pycore_context.h index a833f790a621b1..eb9be65666d0ba 100644 --- a/Include/internal/pycore_context.h +++ b/Include/internal/pycore_context.h @@ -26,13 +26,28 @@ struct _pycontextobject { PyHamtObject *ctx_vars; PyObject *ctx_weakreflist; int ctx_entered; + // Redundant subset of ctx_vars holding exactly the bindings inherited by + // an implicitly created child thread. Since the HAMT is immutable, it can + // be shared directly with the child's context. While this and ctx_vars + // are the same object, one HAMT update serves both fields; adding a + // non-inheritable binding makes them diverge. This relies on both + // var_thread_inherit and thread_inherit_context remaining immutable. + PyHamtObject *ctx_inheritable_vars; }; +typedef enum { + _Py_CONTEXTVAR_INHERIT_DEFAULT = 0, + _Py_CONTEXTVAR_INHERIT_FALSE = 1, + _Py_CONTEXTVAR_INHERIT_TRUE = 2, +} _PyContextVarInherit; + + struct _pycontextvarobject { PyObject_HEAD PyObject *var_name; PyObject *var_default; + _PyContextVarInherit var_thread_inherit; #ifndef Py_GIL_DISABLED PyObject *var_cached; uint64_t var_cached_tsid; @@ -54,6 +69,7 @@ struct _pycontexttokenobject { // _testinternalcapi.hamt() used by tests. // Export for '_testcapi' shared extension PyAPI_FUNC(PyObject*) _PyContext_NewHamtForTests(void); +PyAPI_FUNC(PyObject*) _PyContext_NewThreadContext(void); PyAPI_FUNC(int) _PyContext_Enter(PyThreadState *ts, PyObject *octx); PyAPI_FUNC(int) _PyContext_Exit(PyThreadState *ts, PyObject *octx); diff --git a/Lib/test/test_context.py b/Lib/test/test_context.py index ef20495dcc01ea..ae9cee626a4335 100644 --- a/Lib/test/test_context.py +++ b/Lib/test/test_context.py @@ -9,13 +9,18 @@ import unittest import weakref from test import support -from test.support import threading_helper +from test.support import script_helper, threading_helper try: from _testinternalcapi import hamt except ImportError: hamt = None +try: + import _testcapi +except ImportError: + _testcapi = None + def isolated_context(func): """Needed to make reftracking test mode work.""" @@ -36,12 +41,68 @@ def test_context_var_new_1(self): c = contextvars.ContextVar('aaa') self.assertEqual(c.name, 'aaa') + self.assertIsNone(c.thread_inheritable) with self.assertRaises(AttributeError): c.name = 'bbb' + with self.assertRaises(AttributeError): + c.thread_inheritable = True + + inheritable = contextvars.ContextVar( + 'inheritable', thread_inheritable=True) + self.assertIs(type(inheritable), contextvars.ContextVar) + self.assertEqual(inheritable.name, 'inheritable') + self.assertIs(inheritable.thread_inheritable, True) + + inheritable_with_default = contextvars.ContextVar( + 'inheritable_with_default', default=42, + thread_inheritable=True, + ) + self.assertEqual(inheritable_with_default.get(), 42) + + # None and omission both select the interpreter-wide default. + default_policy = contextvars.ContextVar( + 'default_policy', thread_inheritable=None) + non_inheritable = contextvars.ContextVar( + 'non_inheritable', thread_inheritable=False) + self.assertIsNone(default_policy.thread_inheritable) + self.assertIs(non_inheritable.thread_inheritable, False) + with self.assertRaisesRegex( + TypeError, + 'thread_inheritable must be a bool or None, not int'): + contextvars.ContextVar('var', thread_inheritable=1) + with self.assertRaises(TypeError): + contextvars.ContextVar('var', None) + with self.assertRaises(TypeError): + contextvars.ContextVar('var', inherit=True) + # thread_inheritable is a read-only attribute, not a constructor. + with self.assertRaises(TypeError): + contextvars.ContextVar.thread_inheritable('var') self.assertNotEqual(hash(c), hash('aaa')) + def test_context_var_thread_inheritable_gc(self): + class Value: + pass + + def make_cycle(): + var = contextvars.ContextVar( + 'var', thread_inheritable=True) + ctx = contextvars.Context() + value = Value() + value_ref = weakref.ref(value) + + def bind(): + var.set(value) + value.context = contextvars.copy_context() + + ctx.run(bind) + return value_ref + + value_ref = make_cycle() + support.gc_collect() + self.assertIsNone(value_ref()) + @isolated_context def test_context_var_repr_1(self): c = contextvars.ContextVar('a') @@ -62,6 +123,25 @@ def test_context_var_repr_1(self): c.reset(t) self.assertIn(' used ', repr(t)) + @isolated_context + def test_context_var_repr_thread_inheritable(self): + # The policy is only shown when it overrides the flag. + c = contextvars.ContextVar('a') + self.assertNotIn('thread_inheritable', repr(c)) + c = contextvars.ContextVar('a', thread_inheritable=None) + self.assertNotIn('thread_inheritable', repr(c)) + + c = contextvars.ContextVar('a', thread_inheritable=True) + self.assertRegex( + repr(c), + r"^$") + c = contextvars.ContextVar('a', default=123, thread_inheritable=False) + self.assertRegex( + repr(c), + r"^$") + @isolated_context def test_token_repr_1(self): c = contextvars.ContextVar('a') @@ -587,6 +667,438 @@ def __eq__(self, other): ctx1 == ctx2 +@threading_helper.requires_working_threading() +class ThreadInheritableVarTest(unittest.TestCase): + # These tests run in a subprocess with -X thread_inherit_context pinned, + # since its default depends on the build (true on free-threaded builds). + + def run_with_flag(self, flag, source): + _, _, stderr = script_helper.assert_python_ok( + '-X', f'thread_inherit_context={flag}', '-c', source) + self.assertEqual(stderr, b'') + + def test_thread_inheritance(self): + self.run_with_flag(0, """if True: + import threading + from contextvars import ContextVar, copy_context + + inh = ContextVar('inh', default='default', thread_inheritable=True) + plain = ContextVar('plain') + default_policy = ContextVar( + 'default_policy', thread_inheritable=None) + excluded = ContextVar('excluded', thread_inheritable=False) + inh.set('inherited') + plain.set('not inherited') + default_policy.set('not inherited either') + excluded.set('explicitly excluded') + + def child(): + # The binding is a real binding in the thread's context: + # visible to get(), copy_context() and Context methods. + assert inh.get() == 'inherited' + ctx = copy_context() + assert inh in ctx + assert ctx[inh] == 'inherited' + assert ctx.run(inh.get) == 'inherited' + # Non-inheritable vars are not visible. + assert plain not in ctx + assert default_policy not in ctx + assert excluded not in ctx + for var in (plain, default_policy, excluded): + try: + var.get() + except LookupError: + pass + else: + raise AssertionError(f'{var.name} was inherited') + + t = threading.Thread(target=child) + t.start() + t.join() + """) + + def test_inheritance_captured_at_start_and_context_copy(self): + self.run_with_flag(0, """if True: + import threading + from contextvars import Context, ContextVar + + inh = ContextVar('inh', thread_inheritable=True) + values = [] + + # The binding is captured by start(), not by Thread(). + t = threading.Thread(target=lambda: values.append(inh.get())) + inh.set('at start') + t.start() + t.join() + + def start_and_join(): + t = threading.Thread( + target=lambda: values.append(inh.get())) + t.start() + t.join() + + # Context.run() and Context.copy() preserve the inheritable subset. + ctx = Context() + ctx.run(inh.set, 'context') + ctx.run(start_and_join) + ctx_copy = ctx.copy() + ctx_copy.run(inh.set, 'copy') + ctx_copy.run(start_and_join) + + assert values == ['at start', 'context', 'copy'] + """) + + def test_thread_start_retry_recaptures_context(self): + self.run_with_flag(0, """if True: + import threading + from contextvars import ContextVar + + inh = ContextVar('inh', thread_inheritable=True) + inh.set('first attempt') + values = [] + t = threading.Thread(target=lambda: values.append(inh.get())) + + start_joinable_thread = threading._start_joinable_thread + + def fail_start(*args, **kwargs): + raise threading.ThreadError + + threading._start_joinable_thread = fail_start + try: + try: + t.start() + except threading.ThreadError: + pass + else: + raise AssertionError('thread start did not fail') + finally: + threading._start_joinable_thread = start_joinable_thread + + inh.set('retry') + t.start() + t.join() + assert values == ['retry'] + """) + + def test_thread_set_and_reset(self): + self.run_with_flag(0, """if True: + import threading + from contextvars import ContextVar + + inh = ContextVar('inh', thread_inheritable=True) + inh.set('inherited') + + def child(): + token = inh.set('child value') + assert inh.get() == 'child value' + inh.reset(token) + assert inh.get() == 'inherited' + + t = threading.Thread(target=child) + t.start() + t.join() + # The thread's set() does not affect the parent. + assert inh.get() == 'inherited' + """) + + def test_thread_inheritance_transitive(self): + self.run_with_flag(0, """if True: + import threading + from contextvars import ContextVar + + inh = ContextVar('inh', thread_inheritable=True) + inh.set('inherited') + + def grandchild(): + assert inh.get() == 'inherited' + + def child(): + # The child never sets the var; the binding must still + # propagate to threads it starts. + t = threading.Thread(target=grandchild) + t.start() + t.join() + + t = threading.Thread(target=child) + t.start() + t.join() + """) + + def test_thread_non_inheritable_not_retained(self): + # An excluded binding must be dropped from the new thread's internal + # subset of thread-inheritable bindings, not just from its visible + # bindings. Otherwise the excluded value stays reachable through the + # context of every descendant thread. + source = """if True: + import gc + import threading + import weakref + from contextvars import ContextVar, copy_context + + excluded = ContextVar('excluded', thread_inheritable=False) + descendant = [] + + class Value: + pass + + def chain(depth): + if depth: + t = threading.Thread(target=chain, args=(depth - 1,)) + t.start() + t.join() + else: + descendant.append(copy_context()) + + value = Value() + ref = weakref.ref(value) + excluded.set(value) + chain(4) + del value + excluded.set(None) + gc.collect() + + assert excluded not in descendant[0] + assert ref() is None, 'excluded value was retained by a thread' + """ + for flag in (0, 1): + with self.subTest(thread_inherit_context=flag): + self.run_with_flag(flag, source) + + def test_thread_inheritance_unset_or_deleted(self): + self.run_with_flag(0, """if True: + import threading + from contextvars import ContextVar + + unset = ContextVar('unset', default='default', thread_inheritable=True) + deleted = ContextVar('deleted', thread_inheritable=True) + token = deleted.set('inherited') + deleted.reset(token) + + def child(): + assert unset.get() == 'default' + try: + deleted.get() + except LookupError: + pass + else: + raise AssertionError('deleted binding was inherited') + + t = threading.Thread(target=child) + t.start() + t.join() + """) + + def test_thread_explicit_context(self): + self.run_with_flag(0, """if True: + import threading + from contextvars import ContextVar, Context + + inh = ContextVar('inh', default='default', thread_inheritable=True) + inh.set('inherited') + + def child(): + assert inh.get() == 'default' + + t = threading.Thread(target=child, context=Context()) + t.start() + t.join() + """) + + def test_thread_inheritance_asyncio(self): + self.run_with_flag(0, """if True: + import asyncio + import threading + from contextvars import ContextVar + + inh = ContextVar('inh', thread_inheritable=True) + plain = ContextVar('plain') + inh.set('inherited') + plain.set('not inherited') + + def check_plain_not_set(): + try: + plain.get() + except LookupError: + pass + else: + raise AssertionError('plain was inherited') + + async def task(): + # Tasks run in a copy of the thread's context, which + # includes the inherited binding. + assert inh.get() == 'inherited' + check_plain_not_set() + # A set() inside the task is confined to the task. + inh.set('task value') + + async def main(): + assert inh.get() == 'inherited' + await asyncio.gather(task(), task()) + assert inh.get() == 'inherited' + # Callbacks also run in a copy of the current context. + loop = asyncio.get_running_loop() + fut = loop.create_future() + loop.call_soon( + lambda: fut.set_result((inh.get(), plain.get(None)))) + assert await fut == ('inherited', None) + + def child(): + asyncio.run(main()) + + t = threading.Thread(target=child) + t.start() + t.join() + """) + + def test_thread_inherit_context_flag_true(self): + self.run_with_flag(1, """if True: + import threading + from contextvars import ContextVar, copy_context + + inh = ContextVar('inh', thread_inheritable=True) + plain = ContextVar('plain') + default_policy = ContextVar( + 'default_policy', thread_inheritable=None) + excluded = ContextVar('excluded', thread_inheritable=False) + inh.set('inherited') + plain.set('also inherited') + default_policy.set('also inherited by default') + excluded.set('not inherited') + + def child(): + assert inh.get() == 'inherited' + assert plain.get() == 'also inherited' + assert default_policy.get() == 'also inherited by default' + assert excluded.get(None) is None + assert excluded not in copy_context() + + t = threading.Thread(target=child) + t.start() + t.join() + + # A supplied context is authoritative and is not filtered. + values = [] + t = threading.Thread( + target=lambda: values.append(excluded.get()), + context=copy_context()) + t.start() + t.join() + assert values == ['not inherited'] + """) + + @unittest.skipIf(_testcapi is None, 'requires _testcapi') + def test_capi_thread_inheritance(self): + # The C API flags must override the global flag the same way the + # thread_inheritable keyword argument does. + source = """if True: + import threading + import _testcapi + from contextvars import copy_context + + new_with_flags = _testcapi.contextvar_new_with_flags + always = _testcapi.Py_CONTEXTVAR_INHERIT_THREAD_ALWAYS + never = _testcapi.Py_CONTEXTVAR_INHERIT_THREAD_NEVER + inh = new_with_flags('inh', always) + excl = new_with_flags('excl', never) + plain = _testcapi.contextvar_new('plain') + inh.set('inh value') + excl.set('excl value') + plain.set('plain value') + + values = [] + + def child(): + values.append( + (inh.get(None), excl.get(None), plain.get(None))) + assert excl not in copy_context() + + t = threading.Thread(target=child) + t.start() + t.join() + assert values == [('inh value', None, expected_plain)], values + """ + for flag, expected_plain in ((0, None), (1, 'plain value')): + with self.subTest(thread_inherit_context=flag): + self.run_with_flag( + flag, + f'expected_plain = {expected_plain!r}\n{source}') + + +@unittest.skipIf(_testcapi is None, 'requires _testcapi') +class CAPIContextVarTest(unittest.TestCase): + + def constructors(self): + with_flags = _testcapi.contextvar_new_with_flags + return ( + ('new', _testcapi.contextvar_new), + ('flags_default', lambda name, *default: with_flags( + name, _testcapi.Py_CONTEXTVAR_INHERIT_THREAD_DEFAULT, + *default)), + ('flags_never', lambda name, *default: with_flags( + name, _testcapi.Py_CONTEXTVAR_INHERIT_THREAD_NEVER, + *default)), + ('flags_always', lambda name, *default: with_flags( + name, _testcapi.Py_CONTEXTVAR_INHERIT_THREAD_ALWAYS, + *default)), + ) + + def test_new(self): + var = _testcapi.contextvar_new('plain') + self.assertIs(type(var), contextvars.ContextVar) + self.assertEqual(var.name, 'plain') + self.assertIsNone(var.thread_inheritable) + + def test_new_with_flags(self): + cases = ( + (_testcapi.Py_CONTEXTVAR_INHERIT_THREAD_DEFAULT, None), + (_testcapi.Py_CONTEXTVAR_INHERIT_THREAD_NEVER, False), + (_testcapi.Py_CONTEXTVAR_INHERIT_THREAD_ALWAYS, True), + ) + for flags, expected in cases: + with self.subTest(flags=flags): + var = _testcapi.contextvar_new_with_flags('var', flags) + self.assertIs(type(var), contextvars.ContextVar) + self.assertEqual(var.name, 'var') + self.assertIs(var.thread_inheritable, expected) + + def test_new_with_invalid_flags(self): + always = _testcapi.Py_CONTEXTVAR_INHERIT_THREAD_ALWAYS + never = _testcapi.Py_CONTEXTVAR_INHERIT_THREAD_NEVER + for flags in (-1, always | never, 4): + with self.subTest(flags=flags): + with self.assertRaisesRegex( + ValueError, 'invalid ContextVar flags'): + _testcapi.contextvar_new_with_flags('var', flags) + + def test_no_default(self): + # A NULL default is not the same as a None default. + for name, new in self.constructors(): + with self.subTest(new=name): + var = new('var') + with self.assertRaises(LookupError): + var.get() + self.assertIsNone(var.get(None)) + + def test_default(self): + for name, new in self.constructors(): + with self.subTest(new=name): + self.assertIsNone(new('var', None).get()) + self.assertEqual(new('var', 42).get(), 42) + + def test_default_refcount(self): + # Doc/data/refcounts.dat documents def as a +1 argument. + for name, new in self.constructors(): + with self.subTest(new=name): + default = ['default'] + before = sys.getrefcount(default) + var = new('var', default) + self.assertEqual(sys.getrefcount(default), before + 1) + self.assertIs(var.get(), default) + del var + support.gc_collect() + self.assertEqual(sys.getrefcount(default), before) + + # HAMT Tests diff --git a/Lib/threading.py b/Lib/threading.py index abac31e25886fa..8ede0ecca78c1d 100644 --- a/Lib/threading.py +++ b/Lib/threading.py @@ -1033,8 +1033,8 @@ class is implemented. *context* is the contextvars.Context value to use for the thread. The default value is None, which means to check - sys.flags.thread_inherit_context. If that flag is true, use a copy - of the context of the caller. If false, use an empty context. To + sys.flags.thread_inherit_context. Variables can override that flag + with the ContextVar constructor's thread_inheritable parameter. To explicitly start with an empty context, pass a new instance of contextvars.Context(). To explicitly start with a copy of the current context, pass the value from contextvars.copy_context(). @@ -1127,14 +1127,11 @@ def start(self): with _active_limbo_lock: _limbo[self] = self - if self._context is None: - # No context provided - if _sys.flags.thread_inherit_context: - # start with a copy of the context of the caller - self._context = _contextvars.copy_context() - else: - # start with an empty context - self._context = _contextvars.Context() + context_is_implicit = self._context is None + if context_is_implicit: + # Capture the caller's bindings selected for automatic thread + # inheritance by the global flag and per-variable overrides. + self._context = _contextvars._new_thread_context() try: # Start joinable thread @@ -1143,6 +1140,9 @@ def start(self): except Exception: with _active_limbo_lock: del _limbo[self] + if context_is_implicit: + # Capture the caller's context again if start() is retried. + self._context = None raise self._started.wait() # Will set ident and native_id @@ -1219,6 +1219,8 @@ def _bootstrap_inner(self): except: self._invoke_excepthook(self) finally: + # Break references held by the context after any bootstrap path. + self._context = None self._delete() def _delete(self): diff --git a/Misc/NEWS.d/next/Library/2026-07-23-13-19-38.gh-issue-154562.WBaAJV.rst b/Misc/NEWS.d/next/Library/2026-07-23-13-19-38.gh-issue-154562.WBaAJV.rst new file mode 100644 index 00000000000000..0df0e397fe282d --- /dev/null +++ b/Misc/NEWS.d/next/Library/2026-07-23-13-19-38.gh-issue-154562.WBaAJV.rst @@ -0,0 +1,7 @@ +Add the *thread_inheritable* keyword-only parameter and corresponding +read-only attribute to :class:`contextvars.ContextVar`. They let an +individual variable force its bindings to be included in or excluded from +implicit :class:`threading.Thread` context inheritance, independently of +:data:`sys.flags.thread_inherit_context`. Add +:c:func:`PyContextVar_NewWithFlags` for creating context variables which +override automatic thread context inheritance. diff --git a/Modules/Setup.stdlib.in b/Modules/Setup.stdlib.in index eaf8777a56c59b..a194fe53335524 100644 --- a/Modules/Setup.stdlib.in +++ b/Modules/Setup.stdlib.in @@ -173,7 +173,7 @@ @MODULE__XXTESTFUZZ_TRUE@_xxtestfuzz _xxtestfuzz/_xxtestfuzz.c _xxtestfuzz/fuzzer.c @MODULE__TESTBUFFER_TRUE@_testbuffer _testbuffer.c @MODULE__TESTINTERNALCAPI_TRUE@_testinternalcapi _testinternalcapi.c _testinternalcapi/test_lock.c _testinternalcapi/pytime.c _testinternalcapi/set.c _testinternalcapi/test_critical_sections.c _testinternalcapi/complex.c _testinternalcapi/interpreter.c _testinternalcapi/tuple.c _testinternalcapi/typecache.c -@MODULE__TESTCAPI_TRUE@_testcapi _testcapimodule.c _testcapi/vectorcall.c _testcapi/heaptype.c _testcapi/abstract.c _testcapi/unicode.c _testcapi/dict.c _testcapi/set.c _testcapi/list.c _testcapi/tuple.c _testcapi/getargs.c _testcapi/datetime.c _testcapi/docstring.c _testcapi/mem.c _testcapi/watchers.c _testcapi/long.c _testcapi/float.c _testcapi/complex.c _testcapi/numbers.c _testcapi/structmember.c _testcapi/exceptions.c _testcapi/code.c _testcapi/buffer.c _testcapi/pyatomic.c _testcapi/run.c _testcapi/file.c _testcapi/codec.c _testcapi/immortal.c _testcapi/gc.c _testcapi/hash.c _testcapi/time.c _testcapi/bytes.c _testcapi/object.c _testcapi/modsupport.c _testcapi/monitoring.c _testcapi/config.c _testcapi/import.c _testcapi/frame.c _testcapi/type.c _testcapi/function.c _testcapi/module.c _testcapi/weakref.c +@MODULE__TESTCAPI_TRUE@_testcapi _testcapimodule.c _testcapi/vectorcall.c _testcapi/heaptype.c _testcapi/abstract.c _testcapi/unicode.c _testcapi/dict.c _testcapi/set.c _testcapi/list.c _testcapi/tuple.c _testcapi/getargs.c _testcapi/datetime.c _testcapi/docstring.c _testcapi/mem.c _testcapi/watchers.c _testcapi/long.c _testcapi/float.c _testcapi/complex.c _testcapi/numbers.c _testcapi/structmember.c _testcapi/exceptions.c _testcapi/code.c _testcapi/buffer.c _testcapi/pyatomic.c _testcapi/run.c _testcapi/file.c _testcapi/codec.c _testcapi/context.c _testcapi/immortal.c _testcapi/gc.c _testcapi/hash.c _testcapi/time.c _testcapi/bytes.c _testcapi/object.c _testcapi/modsupport.c _testcapi/monitoring.c _testcapi/config.c _testcapi/import.c _testcapi/frame.c _testcapi/type.c _testcapi/function.c _testcapi/module.c _testcapi/weakref.c @MODULE__TESTLIMITEDCAPI_TRUE@_testlimitedcapi _testlimitedcapi.c _testlimitedcapi/abstract.c _testlimitedcapi/bytearray.c _testlimitedcapi/bytes.c _testlimitedcapi/capsule.c _testlimitedcapi/codec.c _testlimitedcapi/complex.c _testlimitedcapi/dict.c _testlimitedcapi/eval.c _testlimitedcapi/float.c _testlimitedcapi/heaptype_relative.c _testlimitedcapi/import.c _testlimitedcapi/list.c _testlimitedcapi/long.c _testlimitedcapi/object.c _testlimitedcapi/pyos.c _testlimitedcapi/set.c _testlimitedcapi/slots.c _testlimitedcapi/sys.c _testlimitedcapi/threadstate.c _testlimitedcapi/tuple.c _testlimitedcapi/unicode.c _testlimitedcapi/vectorcall_limited.c _testlimitedcapi/version.c _testlimitedcapi/file.c _testlimitedcapi/weakref.c _testlimitedcapi/run.c @MODULE__TESTCLINIC_TRUE@_testclinic _testclinic.c @MODULE__TESTCLINIC_LIMITED_TRUE@_testclinic_limited _testclinic_limited.c diff --git a/Modules/_testcapi/context.c b/Modules/_testcapi/context.c new file mode 100644 index 00000000000000..b56d28bc52c778 --- /dev/null +++ b/Modules/_testcapi/context.c @@ -0,0 +1,62 @@ +#include "parts.h" +#include "util.h" + + +typedef PyObject *(*contextvar_ctor)(const char *, PyObject *); + +static PyObject * +call_contextvar_ctor(contextvar_ctor ctor, PyObject *args) +{ + const char *name; + PyObject *def = NULL; + // Omitting the default passes NULL, which is not the same as Py_None. + if (!PyArg_ParseTuple(args, "s|O", &name, &def)) { + return NULL; + } + return ctor(name, def); +} + +static PyObject * +contextvar_new(PyObject *Py_UNUSED(module), PyObject *args) +{ + return call_contextvar_ctor(PyContextVar_New, args); +} + +static PyObject * +contextvar_new_with_flags(PyObject *Py_UNUSED(module), PyObject *args) +{ + const char *name; + int flags; + PyObject *def = NULL; + // Put flags before the optional default so omitting the default passes + // NULL, which is not the same as Py_None. + if (!PyArg_ParseTuple(args, "si|O", &name, &flags, &def)) { + return NULL; + } + return PyContextVar_NewWithFlags(name, def, flags); +} + + +static PyMethodDef test_methods[] = { + {"contextvar_new", contextvar_new, METH_VARARGS}, + {"contextvar_new_with_flags", contextvar_new_with_flags, METH_VARARGS}, + {NULL}, +}; + +int +_PyTestCapi_Init_Context(PyObject *m) +{ + if (PyModule_AddFunctions(m, test_methods) < 0) { + return -1; + } + if (PyModule_AddIntMacro(m, Py_CONTEXTVAR_INHERIT_THREAD_DEFAULT) < 0) { + return -1; + } + if (PyModule_AddIntMacro(m, Py_CONTEXTVAR_INHERIT_THREAD_NEVER) < 0) { + return -1; + } + if (PyModule_AddIntMacro(m, Py_CONTEXTVAR_INHERIT_THREAD_ALWAYS) < 0) { + return -1; + } + return 0; +} diff --git a/Modules/_testcapi/parts.h b/Modules/_testcapi/parts.h index 98b5dd47accde3..3e3a9f831a81a9 100644 --- a/Modules/_testcapi/parts.h +++ b/Modules/_testcapi/parts.h @@ -68,5 +68,6 @@ int _PyTestCapi_Init_Type(PyObject *mod); int _PyTestCapi_Init_Function(PyObject *mod); int _PyTestCapi_Init_Module(PyObject *mod); int _PyTestCapi_Init_Weakref(PyObject *mod); +int _PyTestCapi_Init_Context(PyObject *mod); #endif // Py_TESTCAPI_PARTS_H diff --git a/Modules/_testcapimodule.c b/Modules/_testcapimodule.c index fb18a866e62812..bf7c408f1d0852 100644 --- a/Modules/_testcapimodule.c +++ b/Modules/_testcapimodule.c @@ -3943,6 +3943,9 @@ _testcapi_exec(PyObject *m) if (_PyTestCapi_Init_Weakref(m) < 0) { return -1; } + if (_PyTestCapi_Init_Context(m) < 0) { + return -1; + } return 0; } diff --git a/PCbuild/_testcapi.vcxproj b/PCbuild/_testcapi.vcxproj index 64e50b67be4656..7caaff38265860 100644 --- a/PCbuild/_testcapi.vcxproj +++ b/PCbuild/_testcapi.vcxproj @@ -134,6 +134,7 @@ + diff --git a/PCbuild/_testcapi.vcxproj.filters b/PCbuild/_testcapi.vcxproj.filters index a3b62e1df663e0..a5a5cc0100edbc 100644 --- a/PCbuild/_testcapi.vcxproj.filters +++ b/PCbuild/_testcapi.vcxproj.filters @@ -135,6 +135,9 @@ Source Files + + Source Files + diff --git a/Python/_contextvars.c b/Python/_contextvars.c index 86acc94fbc79cd..a65e30abf61e03 100644 --- a/Python/_contextvars.c +++ b/Python/_contextvars.c @@ -1,4 +1,5 @@ #include "Python.h" +#include "pycore_context.h" // _PyContext_NewThreadContext() #include "clinic/_contextvars.c.h" @@ -20,10 +21,23 @@ _contextvars_copy_context_impl(PyObject *module) } +/*[clinic input] +_contextvars._new_thread_context +[clinic start generated code]*/ + +static PyObject * +_contextvars__new_thread_context_impl(PyObject *module) +/*[clinic end generated code: output=2d108487bd5c4a49 input=bc928642a19055eb]*/ +{ + return _PyContext_NewThreadContext(); +} + + PyDoc_STRVAR(module_doc, "Context Variables"); static PyMethodDef _contextvars_methods[] = { _CONTEXTVARS_COPY_CONTEXT_METHODDEF + _CONTEXTVARS__NEW_THREAD_CONTEXT_METHODDEF {NULL, NULL} }; diff --git a/Python/clinic/_contextvars.c.h b/Python/clinic/_contextvars.c.h index b1885e41c355d2..dc709782129622 100644 --- a/Python/clinic/_contextvars.c.h +++ b/Python/clinic/_contextvars.c.h @@ -18,4 +18,21 @@ _contextvars_copy_context(PyObject *module, PyObject *Py_UNUSED(ignored)) { return _contextvars_copy_context_impl(module); } -/*[clinic end generated code: output=26e07024451baf52 input=a9049054013a1b77]*/ + +PyDoc_STRVAR(_contextvars__new_thread_context__doc__, +"_new_thread_context($module, /)\n" +"--\n" +"\n"); + +#define _CONTEXTVARS__NEW_THREAD_CONTEXT_METHODDEF \ + {"_new_thread_context", (PyCFunction)_contextvars__new_thread_context, METH_NOARGS, _contextvars__new_thread_context__doc__}, + +static PyObject * +_contextvars__new_thread_context_impl(PyObject *module); + +static PyObject * +_contextvars__new_thread_context(PyObject *module, PyObject *Py_UNUSED(ignored)) +{ + return _contextvars__new_thread_context_impl(module); +} +/*[clinic end generated code: output=c2d927c634f05b08 input=a9049054013a1b77]*/ diff --git a/Python/context.c b/Python/context.c index 4678054ff3ad74..73f6c2562c2e0c 100644 --- a/Python/context.c +++ b/Python/context.c @@ -47,7 +47,8 @@ static PyContext * context_new_empty(void); static PyContext * -context_new_from_vars(PyHamtObject *vars); +context_new_from_vars(PyHamtObject *vars, + PyHamtObject *inheritable_vars); static inline PyContext * context_get(void); @@ -56,7 +57,8 @@ static PyContextToken * token_new(PyContext *ctx, PyContextVar *var, PyObject *val); static PyContextVar * -contextvar_new(PyObject *name, PyObject *def); +contextvar_new(PyObject *name, PyObject *def, + _PyContextVarInherit thread_inherit); static int contextvar_set(PyContextVar *var, PyObject *val); @@ -79,12 +81,30 @@ PyContext_New(void) } +PyObject * +_PyContext_NewThreadContext(void) +{ + PyContext *starter_ctx = context_get(); + if (starter_ctx == NULL) { + return NULL; + } + + // The inheritable subset becomes both the child's full bindings map and + // its inheritable subset. HAMTs are immutable, so both fields can share + // the starter's map without filtering or copying it. + return (PyObject *)context_new_from_vars( + starter_ctx->ctx_inheritable_vars, + starter_ctx->ctx_inheritable_vars); +} + + PyObject * PyContext_Copy(PyObject * octx) { ENSURE_Context(octx, NULL) PyContext *ctx = (PyContext *)octx; - return (PyObject *)context_new_from_vars(ctx->ctx_vars); + return (PyObject *)context_new_from_vars( + ctx->ctx_vars, ctx->ctx_inheritable_vars); } @@ -96,7 +116,8 @@ PyContext_CopyCurrent(void) return NULL; } - return (PyObject *)context_new_from_vars(ctx->ctx_vars); + return (PyObject *)context_new_from_vars( + ctx->ctx_vars, ctx->ctx_inheritable_vars); } static const char * @@ -262,18 +283,49 @@ PyContext_Exit(PyObject *octx) } -PyObject * -PyContextVar_New(const char *name, PyObject *def) +static PyObject * +contextvar_new_from_utf8(const char *name, PyObject *def, + _PyContextVarInherit thread_inherit) { PyObject *pyname = PyUnicode_FromString(name); if (pyname == NULL) { return NULL; } - PyContextVar *var = contextvar_new(pyname, def); + PyContextVar *var = contextvar_new(pyname, def, thread_inherit); Py_DECREF(pyname); return (PyObject *)var; } +PyObject * +PyContextVar_New(const char *name, PyObject *def) +{ + return contextvar_new_from_utf8( + name, def, _Py_CONTEXTVAR_INHERIT_DEFAULT); +} + +PyObject * +PyContextVar_NewWithFlags(const char *name, PyObject *def, int flags) +{ + _PyContextVarInherit thread_inherit; + switch (flags) { + case Py_CONTEXTVAR_INHERIT_THREAD_DEFAULT: + thread_inherit = _Py_CONTEXTVAR_INHERIT_DEFAULT; + break; + case Py_CONTEXTVAR_INHERIT_THREAD_NEVER: + thread_inherit = _Py_CONTEXTVAR_INHERIT_FALSE; + break; + case Py_CONTEXTVAR_INHERIT_THREAD_ALWAYS: + thread_inherit = _Py_CONTEXTVAR_INHERIT_TRUE; + break; + default: + PyErr_Format(PyExc_ValueError, + "invalid ContextVar flags: 0x%x", + (unsigned int)flags); + return NULL; + } + return contextvar_new_from_utf8(name, def, thread_inherit); +} + int PyContextVar_Get(PyObject *ovar, PyObject *def, PyObject **val) @@ -437,6 +489,7 @@ _context_alloc(void) ctx->ctx_vars = NULL; ctx->ctx_prev = NULL; + ctx->ctx_inheritable_vars = NULL; ctx->ctx_entered = 0; ctx->ctx_weakreflist = NULL; @@ -458,13 +511,20 @@ context_new_empty(void) return NULL; } + // A fresh context has no excluded bindings, so its full and inheritable + // maps are the same object. Preserve this identity to let one HAMT update + // serve both fields until a non-inheritable binding is added. + ctx->ctx_inheritable_vars = + (PyHamtObject *)Py_NewRef(ctx->ctx_vars); + _PyObject_GC_TRACK(ctx); return ctx; } static PyContext * -context_new_from_vars(PyHamtObject *vars) +context_new_from_vars(PyHamtObject *vars, + PyHamtObject *inheritable_vars) { PyContext *ctx = _context_alloc(); if (ctx == NULL) { @@ -472,6 +532,8 @@ context_new_from_vars(PyHamtObject *vars) } ctx->ctx_vars = (PyHamtObject*)Py_NewRef(vars); + ctx->ctx_inheritable_vars = + (PyHamtObject*)Py_NewRef(inheritable_vars); _PyObject_GC_TRACK(ctx); return ctx; @@ -523,6 +585,7 @@ context_tp_clear(PyObject *op) PyContext *self = _PyContext_CAST(op); Py_CLEAR(self->ctx_prev); Py_CLEAR(self->ctx_vars); + Py_CLEAR(self->ctx_inheritable_vars); return 0; } @@ -532,6 +595,7 @@ context_tp_traverse(PyObject *op, visitproc visit, void *arg) PyContext *self = _PyContext_CAST(op); Py_VISIT(self->ctx_prev); Py_VISIT(self->ctx_vars); + Py_VISIT(self->ctx_inheritable_vars); return 0; } @@ -708,7 +772,8 @@ static PyObject * _contextvars_Context_copy_impl(PyContext *self) /*[clinic end generated code: output=30ba8896c4707a15 input=ebafdbdd9c72d592]*/ { - return (PyObject *)context_new_from_vars(self->ctx_vars); + return (PyObject *)context_new_from_vars( + self->ctx_vars, self->ctx_inheritable_vars); } @@ -782,6 +847,22 @@ PyTypeObject PyContext_Type = { /////////////////////////// ContextVar +static int +contextvar_is_thread_inheritable(PyContextVar *var) +{ + switch (var->var_thread_inherit) { + case _Py_CONTEXTVAR_INHERIT_DEFAULT: + return _PyInterpreterState_GET()->config.thread_inherit_context; + case _Py_CONTEXTVAR_INHERIT_FALSE: + return 0; + case _Py_CONTEXTVAR_INHERIT_TRUE: + return 1; + } + // No default case so the compiler warns about unhandled policies. + Py_UNREACHABLE(); +} + + static int contextvar_set(PyContextVar *var, PyObject *val) { @@ -795,19 +876,49 @@ contextvar_set(PyContextVar *var, PyObject *val) return -1; } + PyHamtObject *old_vars = ctx->ctx_vars; PyHamtObject *new_vars = _PyHamt_Assoc( - ctx->ctx_vars, (PyObject *)var, val); + old_vars, (PyObject *)var, val); if (new_vars == NULL) { return -1; } - Py_SETREF(ctx->ctx_vars, new_vars); + // Compute both new maps before installing either so that an error + // leaves the context unchanged. + PyHamtObject *old_inheritable_vars = NULL; + PyHamtObject *new_inheritable_vars = NULL; + if (contextvar_is_thread_inheritable(var)) { + old_inheritable_vars = ctx->ctx_inheritable_vars; + if (old_inheritable_vars == old_vars) { + // All current bindings are inheritable. One update serves both + // fields and preserves sharing for subsequent updates. + new_inheritable_vars = + (PyHamtObject *)Py_NewRef(new_vars); + } + else { + new_inheritable_vars = _PyHamt_Assoc( + old_inheritable_vars, (PyObject *)var, val); + if (new_inheritable_vars == NULL) { + Py_DECREF(new_vars); + return -1; + } + } + } + // Install both maps before either old map is decref'ed. Besides keeping + // their subset relationship atomic with respect to finalizers, updating + // the cache first lets a re-entrant set() performed by a finalizer win. + ctx->ctx_vars = new_vars; + if (new_inheritable_vars != NULL) { + ctx->ctx_inheritable_vars = new_inheritable_vars; + } #ifndef Py_GIL_DISABLED var->var_cached = val; /* borrow */ var->var_cached_tsid = ts->id; var->var_cached_tsver = ts->context_ver; #endif + Py_DECREF(old_vars); + Py_XDECREF(old_inheritable_vars); return 0; } @@ -823,19 +934,48 @@ contextvar_del(PyContextVar *var) return -1; } - PyHamtObject *vars = ctx->ctx_vars; - PyHamtObject *new_vars = _PyHamt_Without(vars, (PyObject *)var); + PyHamtObject *old_vars = ctx->ctx_vars; + PyHamtObject *new_vars = _PyHamt_Without(old_vars, (PyObject *)var); if (new_vars == NULL) { return -1; } - if (vars == new_vars) { + if (old_vars == new_vars) { Py_DECREF(new_vars); PyErr_SetObject(PyExc_LookupError, (PyObject *)var); return -1; } - Py_SETREF(ctx->ctx_vars, new_vars); + // Compute both new maps before installing either so that an error + // leaves the context unchanged. + PyHamtObject *old_inheritable_vars = NULL; + PyHamtObject *new_inheritable_vars = NULL; + if (contextvar_is_thread_inheritable(var)) { + old_inheritable_vars = ctx->ctx_inheritable_vars; + if (old_inheritable_vars == old_vars) { + // Both fields contain the same bindings, so the result of the + // full-map deletion is also the new inheritable map. + new_inheritable_vars = + (PyHamtObject *)Py_NewRef(new_vars); + } + else { + new_inheritable_vars = _PyHamt_Without( + old_inheritable_vars, (PyObject *)var); + if (new_inheritable_vars == NULL) { + Py_DECREF(new_vars); + return -1; + } + } + } + + // Keep the full map and its inheritable subset in sync before decrefing + // either displaced map. + ctx->ctx_vars = new_vars; + if (new_inheritable_vars != NULL) { + ctx->ctx_inheritable_vars = new_inheritable_vars; + } + Py_DECREF(old_vars); + Py_XDECREF(old_inheritable_vars); return 0; } @@ -868,7 +1008,8 @@ contextvar_generate_hash(void *addr, PyObject *name) } static PyContextVar * -contextvar_new(PyObject *name, PyObject *def) +contextvar_new(PyObject *name, PyObject *def, + _PyContextVarInherit thread_inherit) { if (!PyUnicode_Check(name)) { PyErr_SetString(PyExc_TypeError, @@ -883,6 +1024,7 @@ contextvar_new(PyObject *name, PyObject *def) var->var_name = Py_NewRef(name); var->var_default = Py_XNewRef(def); + var->var_thread_inherit = thread_inherit; #ifndef Py_GIL_DISABLED var->var_cached = NULL; @@ -917,17 +1059,43 @@ class _contextvars.ContextVar "PyContextVar *" "&PyContextVar_Type" static PyObject * contextvar_tp_new(PyTypeObject *type, PyObject *args, PyObject *kwds) { - static char *kwlist[] = {"", "default", NULL}; + static char *kwlist[] = { + "", "default", "thread_inheritable", NULL, + }; PyObject *name; PyObject *def = NULL; + PyObject *thread_inheritable = NULL; if (!PyArg_ParseTupleAndKeywords( - args, kwds, "O|$O:ContextVar", kwlist, &name, &def)) + args, kwds, "O|$OO:ContextVar", kwlist, + &name, &def, &thread_inheritable)) { return NULL; } - return (PyObject *)contextvar_new(name, def); + if (thread_inheritable != NULL && + thread_inheritable != Py_None && + thread_inheritable != Py_True && + thread_inheritable != Py_False) { + PyErr_Format( + PyExc_TypeError, + "thread_inheritable must be a bool or None, not %T", + thread_inheritable); + return NULL; + } + + _PyContextVarInherit thread_inherit; + if (thread_inheritable == Py_True) { + thread_inherit = _Py_CONTEXTVAR_INHERIT_TRUE; + } + else if (thread_inheritable == Py_False) { + thread_inherit = _Py_CONTEXTVAR_INHERIT_FALSE; + } + else { + thread_inherit = _Py_CONTEXTVAR_INHERIT_DEFAULT; + } + + return (PyObject *)contextvar_new(name, def, thread_inherit); } static int @@ -972,11 +1140,22 @@ static PyObject * contextvar_tp_repr(PyObject *op) { PyContextVar *self = _PyContextVar_CAST(op); + const char *thread_inherit_repr = NULL; + if (self->var_thread_inherit == _Py_CONTEXTVAR_INHERIT_TRUE) { + thread_inherit_repr = " thread_inheritable=True"; + } + else if (self->var_thread_inherit == _Py_CONTEXTVAR_INHERIT_FALSE) { + thread_inherit_repr = " thread_inheritable=False"; + } + Py_ssize_t thread_inherit_repr_len = thread_inherit_repr == NULL + ? 0 : (Py_ssize_t)strlen(thread_inherit_repr); + // Estimation based on the shortest name and default value, // but maximize the pointer size. // "" // "" Py_ssize_t estimate = self->var_default ? 53 : 43; + estimate += thread_inherit_repr_len; PyUnicodeWriter *writer = PyUnicodeWriter_Create(estimate); if (writer == NULL) { return NULL; @@ -998,6 +1177,14 @@ contextvar_tp_repr(PyObject *op) } } + // Only shown when set, so that the common case is unchanged. + if (thread_inherit_repr != NULL && + PyUnicodeWriter_WriteASCII( + writer, thread_inherit_repr, thread_inherit_repr_len) < 0) + { + goto error; + } + if (PyUnicodeWriter_Format(writer, " at %p>", self) < 0) { goto error; } @@ -1093,11 +1280,33 @@ _contextvars_ContextVar_reset_impl(PyContextVar *self, PyObject *token) } +static PyObject * +contextvar_get_thread_inheritable(PyObject *op, void *Py_UNUSED(ignored)) +{ + PyContextVar *self = _PyContextVar_CAST(op); + switch (self->var_thread_inherit) { + case _Py_CONTEXTVAR_INHERIT_DEFAULT: + Py_RETURN_NONE; + case _Py_CONTEXTVAR_INHERIT_FALSE: + Py_RETURN_FALSE; + case _Py_CONTEXTVAR_INHERIT_TRUE: + Py_RETURN_TRUE; + } + // No default case so the compiler warns about unhandled policies. + Py_UNREACHABLE(); +} + static PyMemberDef PyContextVar_members[] = { {"name", _Py_T_OBJECT, offsetof(PyContextVar, var_name), Py_READONLY}, {NULL} }; +static PyGetSetDef PyContextVar_getsetlist[] = { + {"thread_inheritable", contextvar_get_thread_inheritable, NULL, + PyDoc_STR("the automatic thread-inheritance policy: None, True, or False")}, + {NULL} +}; + static PyMethodDef PyContextVar_methods[] = { _CONTEXTVARS_CONTEXTVAR_GET_METHODDEF _CONTEXTVARS_CONTEXTVAR_SET_METHODDEF @@ -1114,6 +1323,7 @@ PyTypeObject PyContextVar_Type = { sizeof(PyContextVar), .tp_methods = PyContextVar_methods, .tp_members = PyContextVar_members, + .tp_getset = PyContextVar_getsetlist, .tp_dealloc = contextvar_tp_dealloc, .tp_getattro = PyObject_GenericGetAttr, .tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,