Add the concept of a tree variable (TreeVar)

docs/source/reference-core.rst

@@ -656,6 +656,8 @@ finished.
 
    This code will wait 5 seconds (for the child task to finish), and then return.
 
+.. _child-tasks-and-cancellation:
+
 Child tasks and cancellation
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
@@ -933,8 +935,8 @@ Will properly log the inner exceptions:
 
 .. _task-local-storage:
 
-Task-local storage
-------------------
+Context variables support task-local storage
+--------------------------------------------
 
 Suppose you're writing a server that responds to network requests, and
 you log some information about each request as you process it. If the
@@ -1020,6 +1022,94 @@ For more information, read the
 `contextvar docs <https://docs.python.org/3.7/library/contextvars.html>`__.
 
 
+More on context variable inheritance
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+As you can see from the example above, a task's context variables are
+automatically inherited by any other tasks it starts. To be precise,
+each new task gets a *shallow copy* of the context in the task that
+spawned it. That means:
+
+* If the new task changes one of its context variables using `ContextVar.set()
+  <contextvars.ContextVar.set>`, then that change is not visible in the
+  task that started it. (This would hardly be "task-local" storage otherwise!)
+
+* But if the context variable referred to a mutable object (such as a
+  list or dictionary), and the new task makes a change to that object
+  (such as by calling ``some_contextvar.get().append(42)``), then that
+  change *is* visible in the task that started it, as well as in any other
+  tasks it started. Since that's rather confusing, it's often best to
+  limit yourself to immutable values (strings, integers, tuples, and so
+  on) when working with context variables.
+
+A new task's context is set up as a copy of the context that existed
+at the time of the call to :meth:`~Nursery.start_soon` or
+:meth:`~Nursery.start` that created it.  For example, this code:
+
+.. code-block:: python3
+
+   some_cvar = contextvars.ContextVar()
+
+   async def print_in_child(tag):
+       print("In child", tag, "some_cvar has value", some_cvar.get())
+
+   some_cvar.set(1)
+   async with trio.open_nursery() as nursery:
+       nursery.start_soon(print_in_child, 1)
+       some_cvar.set(2)
+       nursery.start_soon(print_in_child, 2)
+       some_cvar.set(3)
+       print("In parent some_cvar has value", some_cvar.get())
+
+will produce output like::
+
+    In parent some_cvar has value 3
+    In child 1 some_cvar has value 1
+    In child 2 some_cvar has value 2
+
+(If you run it yourself, you might find that the "child 2" line comes
+before "child 1", but it will still be the case that child 1 sees value 1
+while child 2 sees value 2.)
+
+You might wonder why this differs from the behavior of cancel scopes,
+which only apply to a new task if they surround the new task's entire
+nursery (as explained above in the section on
+:ref:`child-tasks-and-cancellation`). The difference is that a cancel
+scope has a limited lifetime (it can't cancel anything once you exit
+its ``with`` block), while a context variable's value is just a value
+(request #42 can keep being request #42 for as long as it likes,
+without any cooperation from the task that created it).
+
+In specialized cases, you might want to provide a task-local value
+that's inherited only from the parent nursery, like cancel scopes are.
+(For example, maybe you're trying to provide child tasks with access
+to a limited-lifetime resource such as a nursery or network
+connection, and you only want a task to be able to use the resource if
+it's going to remain available for the task's entire lifetime.)  Trio
+supports this using `TreeVar`, which is like `contextvars.ContextVar`
+except for the way that it's inherited by new tasks. (It's a "tree"
+variable because it's inherited along the parent-child links that
+form the Trio task tree.)
+
+If the above example used `TreeVar`, then its output would be:
+
+.. code-block:: none
+   :emphasize-lines: 3
+
+   In parent some_cvar has value 3
+   In child 1 some_cvar has value 1
+   In child 2 some_cvar has value 1
+
+because child 2 would inherit the value from its parent nursery, rather than
+from the environment of the ``start_soon()`` call that creates it.
+
+.. autoclass:: trio.TreeVar(name, [*, default])
+
+   .. automethod:: being
+      :with:
+   .. automethod:: get_in(task_or_nursery, [default])
+
+
 .. _synchronization:
 
 Synchronizing and communicating between tasks

docs/source/reference-lowlevel.rst

@@ -256,7 +256,11 @@ anything real. See `#26
 Global state: system tasks and run-local variables
 ==================================================
 
-.. autoclass:: RunVar
+.. autoclass:: RunVar(name, [default])
+
+   .. automethod:: get([default])
+   .. automethod:: set
+   .. automethod:: reset
 
 .. autofunction:: spawn_system_task
 

newsfragments/1523.feature.rst

@@ -0,0 +1,8 @@
+Added the concept of a "tree variable" (`trio.TreeVar`), which is like
+a context variable except that its value in a new task is inherited
+from the new task's parent nursery rather than from the task that
+spawned it.  (The `~trio.TreeVar` behavior matches the existing
+:ref:`behavior of cancel scopes <child-tasks-and-cancellation>`.)
+This distinction makes tree variables useful for anything that's
+naturally inherited along parent/child task relationships, such as a
+reference to a resource that has a limited lifetime.

trio/__init__.py

@@ -32,6 +32,7 @@
     BrokenResourceError,
     EndOfChannel,
     Nursery,
+    TreeVar,
 )
 
 from ._timeouts import (

trio/_core/__init__.py

@@ -69,7 +69,7 @@
 
 from ._unbounded_queue import UnboundedQueue
 
-from ._local import RunVar
+from ._local import RunVar, TreeVar
 
 from ._thread_cache import start_thread_soon
 

trio/_core/_local.py

@@ -1,7 +1,9 @@
-# Runvar implementations
+# Implementations of RunVar and ScopeVar
+import contextvars
+from contextlib import contextmanager
 from . import _run
 
-from .._util import SubclassingDeprecatedIn_v0_15_0
+from .._util import Final, SubclassingDeprecatedIn_v0_15_0
 
 
 class _RunVarToken:
@@ -19,42 +21,46 @@ def __init__(self, var, value):
         self.redeemed = False
 
 
+_NO_DEFAULT = object()
+
+
 class RunVar(metaclass=SubclassingDeprecatedIn_v0_15_0):
     """The run-local variant of a context variable.
 
     :class:`RunVar` objects are similar to context variable objects,
     except that they are shared across a single call to :func:`trio.run`
     rather than a single task.
-
     """
 
-    _NO_DEFAULT = object()
     __slots__ = ("_name", "_default")
 
     def __init__(self, name, default=_NO_DEFAULT):
         self._name = name
         self._default = default
 
     def get(self, default=_NO_DEFAULT):
-        """Gets the value of this :class:`RunVar` for the current run call."""
+        """Gets the value of this `RunVar` for the current call
+        to :func:`trio.run`."""
         try:
             return _run.GLOBAL_RUN_CONTEXT.runner._locals[self]
         except AttributeError:
             raise RuntimeError("Cannot be used outside of a run context") from None
         except KeyError:
             # contextvars consistency
-            if default is not self._NO_DEFAULT:
+            if default is not _NO_DEFAULT:
                 return default
 
-            if self._default is not self._NO_DEFAULT:
+            if self._default is not _NO_DEFAULT:
                 return self._default
 
             raise LookupError(self) from None
 
     def set(self, value):
-        """Sets the value of this :class:`RunVar` for this current run
-        call.
+        """Sets the value of this `RunVar` for the current call to
+        :func:`trio.run`.
 
+        Returns a token which may be passed to :meth:`reset` to restore
+        the previous value.
         """
         try:
             old_value = self.get()
@@ -69,9 +75,8 @@ def set(self, value):
         return token
 
     def reset(self, token):
-        """Resets the value of this :class:`RunVar` to what it was
-        previously specified by the token.
-
+        """Resets the value of this `RunVar` to the value it had
+        before the call to :meth:`set` that returned the given *token*.
         """
         if token is None:
             raise TypeError("token must not be none")
@@ -95,3 +100,118 @@ def reset(self, token):
 
     def __repr__(self):
         return "<RunVar name={!r}>".format(self._name)
+
+
+class TreeVar(metaclass=Final):
+    """A "tree variable": like a context variable except that its value
+    in a new task is inherited from the new task's parent nursery rather
+    than from the new task's spawner.
+
+    `TreeVar` objects support all the same methods and attributes as
+    `~contextvars.ContextVar` objects
+    (:meth:`~contextvars.ContextVar.get`,
+    :meth:`~contextvars.ContextVar.set`,
+    :meth:`~contextvars.ContextVar.reset`, and
+    `~contextvars.ContextVar.name`), and they are constructed the same
+    way. They also provide the additional methods :meth:`being` and
+    :meth:`get_in`, documented below.
+
+    Accessing or changing the value of a `TreeVar` outside of a Trio
+    task will raise `RuntimeError`. (Exception: :meth:`get_in` still
+    works outside of a task, as long as you have a reference to the
+    task or nursery of interest.)
+
+    .. note:: `TreeVar` values are not directly stored in the
+       `contextvars.Context`, so you can't use `Context.get()
+       <contextvars.Context.get>` to access them. If you need the value
+       in a context other than your own, use :meth:`get_in`.
+
+    """
+
+    __slots__ = ("_cvar",)
+
+    def __init__(self, name, **default):
+        self._cvar = contextvars.ContextVar(name, **default)
+
+    @property
+    def name(self):
+        """The name of the variable, as passed during construction. Read-only."""
+        return self._cvar.name
+
+    def get(self, default=_NO_DEFAULT):
+        """Gets the value of this `TreeVar` for the current task.
+
+        If this `TreeVar` has no value in the current task, then
+        :meth:`get` returns the *default* specified as argument to
+        :meth:`get`, or else the *default* specified when constructing
+        the `TreeVar`, or else raises `LookupError`. See the
+        documentation of :meth:`contextvars.ContextVar.get` for more
+        details.
+        """
+        # This is effectively an inlining for efficiency of:
+        # return _run.current_task()._tree_context.run(self._cvar.get, default)
+        try:
+            return _run.GLOBAL_RUN_CONTEXT.task._tree_context[self._cvar]
+        except AttributeError:
+            raise RuntimeError("must be called from async context") from None
+        except KeyError:
+            pass
+        # This will always return the default or raise, because we never give
+        # self._cvar a value in any context in which we run user code.
+        if default is _NO_DEFAULT:
+            return self._cvar.get()
+        else:
+            return self._cvar.get(default)
+
+    def set(self, value):
+        """Sets the value of this `TreeVar` for the current task.  The new
+        value will be inherited by nurseries that are later opened in
+        this task, so that new tasks can inherit whatever value was
+        set when their parent nursery was created.
+
+        Returns a token which may be passed to :meth:`reset` to restore
+        the previous value.
+        """
+        return _run.current_task()._tree_context.run(self._cvar.set, value)
+
+    def reset(self, token):
+        """Resets the value of this `TreeVar` to the value it had
+        before the call to :meth:`set` that returned the given *token*.
+
+        The *token* must have been obtained from a call to :meth:`set` on
+        this same `TreeVar` and in the same task that is now calling
+        :meth:`reset`. Also, each *token* may only be used in one call to
+        :meth:`reset`. Violating these conditions will raise `ValueError`.
+        """
+        _run.current_task()._tree_context.run(self._cvar.reset, token)
+
+    @contextmanager
+    def being(self, value):
+        """Returns a context manager which sets the value of this `TreeVar` to
+        *value* upon entry and restores its previous value upon exit.
+        """
+        token = self.set(value)
+        try:
+            yield
+        finally:
+            self.reset(token)
+
+    def get_in(self, task_or_nursery, default=_NO_DEFAULT):
+        """Gets the value of this `TreeVar` in the given
+        `~trio.lowlevel.Task` or `~trio.Nursery`.
+
+        The value in a task is the value that would be returned by a
+        call to :meth:`~contextvars.ContextVar.get` in that task. The
+        value in a nursery is the value that would be returned by
+        :meth:`~contextvars.ContextVar.get` at the beginning of a new
+        child task started in that nursery. The *default* argument has
+        the same semantics as it does for :meth:`~contextvars.ContextVar.get`.
+        """
+        # copy() so this works from a different thread too. It's a
+        # cheap and thread-safe operation (just copying one reference)
+        # since the underlying context data is immutable.
+        defarg = () if default is _NO_DEFAULT else (default,)
+        return task_or_nursery._tree_context.copy().run(self._cvar.get, *defarg)
+
+    def __repr__(self):
+        return f"<TreeVar name={self.name!r}>"