Add TreeVar

A TreeVar acts like a context variable that is inherited at nursery creation time (and then by child tasks of that nursery) rather than at task creation time. They are useful for providing 'ambient' access to a resource that is tied to an `async with` block in the parent task, such as an open file or trio-asyncio event loop.

Prior art: python-trio/trio#1543 (never made it into mainline Trio). The implementation without Trio core support is somewhat less efficient, but still workable.

.coveragerc

@@ -8,3 +8,4 @@ exclude_lines =
   pragma: no cover
   abc.abstractmethod
   if TYPE_CHECKING:
+  @overload

.github/workflows/ci.yml

@@ -40,6 +40,11 @@ jobs:
         env:
           # Should match 'name:' up above
           JOB_NAME: 'Windows (${{ matrix.python }}, ${{ matrix.arch }}${{ matrix.extra_name }})'
+      - uses: codecov/codecov-action@v3
+        with:
+          directory: empty
+          name: 'Windows (${{ matrix.python }}, ${{ matrix.arch }})'
+          flags: Windows,${{ matrix.python }}
 
   Ubuntu:
     name: 'Ubuntu (${{ matrix.python }}${{ matrix.extra_name }})'
@@ -85,6 +90,11 @@ jobs:
           CHECK_LINT: '${{ matrix.check_lint }}'
           # Should match 'name:' up above
           JOB_NAME: 'Ubuntu (${{ matrix.python }}${{ matrix.extra_name }})'
+      - uses: codecov/codecov-action@v3
+        with:
+          directory: empty
+          name: 'Ubuntu (${{ matrix.python }}${{ matrix.extra_name }})'
+          flags: Ubuntu,${{ matrix.python }}
 
   macOS:
     name: 'macOS (${{ matrix.python }})'
@@ -113,3 +123,8 @@ jobs:
         env:
           # Should match 'name:' up above
           JOB_NAME: 'macOS (${{ matrix.python }})'
+      - uses: codecov/codecov-action@v3
+        with:
+          directory: empty
+          name: 'macOS (${{ matrix.python }})'
+          flags: macOS,${{ matrix.python }}

README.rst

@@ -5,14 +5,14 @@ tricycle: experimental extensions for Trio
    :target: https://pypi.org/project/tricycle
    :alt: Latest PyPI version
 
+.. image:: https://github.com/oremanj/tricycle/actions/workflows/ci.yml/badge.svg
+   :target: https://github.com/oremanj/tricycle/actions/workflows/ci.yml
+   :alt: Automated test status
+
 .. image:: https://img.shields.io/badge/docs-read%20now-blue.svg
    :target: https://tricycle.readthedocs.io/en/latest/?badge=latest
    :alt: Documentation status
 
-.. image:: https://travis-ci.org/oremanj/tricycle.svg?branch=master
-   :target: https://travis-ci.org/oremanj/tricycle
-   :alt: Automated test status
-
 .. image:: https://codecov.io/gh/oremanj/tricycle/branch/master/graph/badge.svg
    :target: https://codecov.io/gh/oremanj/tricycle
    :alt: Test coverage

docs/source/conf.py

@@ -22,6 +22,20 @@
 # So autodoc can import our package
 sys.path.insert(0, os.path.abspath('../..'))
 
+# https://docs.readthedocs.io/en/stable/builds.html#build-environment
+if "READTHEDOCS" in os.environ:
+    import glob
+
+    if glob.glob("../../newsfragments/*.*.rst"):
+        print("-- Found newsfragments; running towncrier --", flush=True)
+        import subprocess
+
+        subprocess.run(
+            ["towncrier", "--yes", "--date", "not released yet"],
+            cwd="../..",
+            check=True,
+        )
+
 # Warn about all references to unknown targets
 nitpicky = True
 # Except for these ones, which we expect to point to unknown targets:

docs/source/reference.rst

@@ -214,3 +214,78 @@ from above would reduce to::
     async def use_websocket():
         async with WebsocketConnection(**etc) as conn:
             await conn.send("Hi!")
+
+
+.. _tree-variables:
+
+Tree variables
+--------------
+
+When you start a new Trio task, the initial values of its `context variables
+<https://trio.readthedocs.io/en/stable/reference-core.html#task-local-storage>`__
+(`contextlib.ContextVar`) are inherited from the environment of the
+`~trio.Nursery.start_soon` or `~trio.Nursery.start` call that
+started the new task. 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 in the Trio documentation about
+`child tasks and cancellation <https://trio.readthedocs.io/en/stable/reference-core.html#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. You can support
+this use case 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:: tricycle.TreeVar(name, [*, default])
+
+   .. automethod:: being
+      :with:
+   .. automethod:: get_in(task_or_nursery, [default])

newsfragments/18.feature.rst

@@ -0,0 +1,6 @@
+Added `tricycle.TreeVar`, which acts like a context variable that is
+inherited at nursery creation time (and then by child tasks of that
+nursery) rather than at task creation time. :ref:`Tree variables
+<tree-variables>` are useful for providing safe 'ambient' access to a
+resource that is tied to an `async with` block in the parent task,
+such as an open file or trio-asyncio event loop.

tricycle/__init__.py

@@ -8,6 +8,7 @@
 from ._multi_cancel import MultiCancelScope as MultiCancelScope
 from ._service_nursery import open_service_nursery as open_service_nursery
 from ._meta import ScopedObject as ScopedObject, BackgroundObject as BackgroundObject
+from ._tree_var import TreeVar as TreeVar, TreeVarToken as TreeVarToken
 
 # watch this space...
 

tricycle/_tests/test_tree_var.py

@@ -0,0 +1,144 @@
+import pytest
+import trio
+import trio.testing
+from functools import partial
+from trio_typing import TaskStatus
+from typing import Optional, Any, cast
+
+from .. import TreeVar, TreeVarToken
+
+
+async def test_treevar() -> None:
+    tv1 = TreeVar[int]("tv1")
+    tv2 = TreeVar[Optional[int]]("tv2", default=None)
+    tv3 = TreeVar("tv3", default=-1)
+    assert tv1.name == "tv1"
+    assert "TreeVar name='tv2'" in repr(tv2)
+
+    with pytest.raises(LookupError):
+        tv1.get()
+    assert tv2.get() is None
+    assert tv1.get(42) == 42
+    assert tv2.get(42) == 42
+
+    NOTHING = cast(int, object())
+
+    async def should_be(val1: int, val2: int, new1: int = NOTHING) -> None:
+        assert tv1.get(NOTHING) == val1
+        assert tv2.get(NOTHING) == val2
+        if new1 is not NOTHING:
+            tv1.set(new1)
+
+    tok1 = tv1.set(10)
+    async with trio.open_nursery() as outer:
+        tok2 = tv1.set(15)
+        with tv2.being(20):
+            assert tv2.get_in(trio.lowlevel.current_task()) == 20
+            async with trio.open_nursery() as inner:
+                tv1.reset(tok2)
+                outer.start_soon(should_be, 10, NOTHING, 100)
+                inner.start_soon(should_be, 15, 20, 200)
+                await trio.testing.wait_all_tasks_blocked()
+                assert tv1.get_in(trio.lowlevel.current_task()) == 10
+                await should_be(10, 20, 300)
+                assert tv1.get_in(inner) == 15
+                assert tv1.get_in(outer) == 10
+                assert tv1.get_in(trio.lowlevel.current_task()) == 300
+                assert tv2.get_in(inner) == 20
+                assert tv2.get_in(outer) is None
+                assert tv2.get_in(trio.lowlevel.current_task()) == 20
+                tv1.reset(tok1)
+                await should_be(NOTHING, 20)
+                assert tv1.get_in(inner) == 15
+                assert tv1.get_in(outer) == 10
+                with pytest.raises(LookupError):
+                    assert tv1.get_in(trio.lowlevel.current_task())
+                # Test get_in() needing to search a parent task but
+                # finding no value there:
+                tv3 = TreeVar("tv3", default=-1)
+                assert tv3.get_in(outer) == -1
+                assert tv3.get_in(outer, -42) == -42
+        assert tv2.get() is None
+        assert tv2.get_in(trio.lowlevel.current_task()) is None
+
+
+def trivial_abort(_: object) -> trio.lowlevel.Abort:
+    return trio.lowlevel.Abort.SUCCEEDED  # pragma: no cover
+
+
+async def test_treevar_follows_eventual_parent():
+    tv1 = TreeVar[str]("tv1")
+
+    async def manage_target(task_status: TaskStatus[trio.Nursery]) -> None:
+        assert tv1.get() == "source nursery"
+        with tv1.being("target nursery"):
+            assert tv1.get() == "target nursery"
+            async with trio.open_nursery() as target_nursery:
+                with tv1.being("target nested child"):
+                    assert tv1.get() == "target nested child"
+                    task_status.started(target_nursery)
+                    await trio.lowlevel.wait_task_rescheduled(trivial_abort)
+                    assert tv1.get() == "target nested child"
+                assert tv1.get() == "target nursery"
+            assert tv1.get() == "target nursery"
+        assert tv1.get() == "source nursery"
+
+    async def verify(
+        value: str, *, task_status: TaskStatus[None] = trio.TASK_STATUS_IGNORED
+    ) -> None:
+        assert tv1.get() == value
+        task_status.started()
+        assert tv1.get() == value
+
+    with tv1.being("source nursery"):
+        async with trio.open_nursery() as source_nursery:
+            with tv1.being("source->target start call"):
+                target_nursery = await source_nursery.start(manage_target)
+            with tv1.being("verify task"):
+                source_nursery.start_soon(verify, "source nursery")
+                target_nursery.start_soon(verify, "target nursery")
+                await source_nursery.start(verify, "source nursery")
+                await target_nursery.start(verify, "target nursery")
+            trio.lowlevel.reschedule(target_nursery.parent_task)
+
+
+async def test_treevar_token_bound_to_task_that_obtained_it() -> None:
+    tv1 = TreeVar[int]("tv1")
+    token: Optional[TreeVarToken[int]] = None
+
+    async def get_token():
+        nonlocal token
+        token = tv1.set(10)
+        try:
+            await trio.lowlevel.wait_task_rescheduled(trivial_abort)
+        finally:
+            tv1.reset(token)
+            with pytest.raises(LookupError):
+                tv1.get()
+            with pytest.raises(LookupError):
+                tv1.get_in(trio.lowlevel.current_task())
+
+    async with trio.open_nursery() as nursery:
+        nursery.start_soon(get_token)
+        await trio.testing.wait_all_tasks_blocked()
+        assert token is not None
+        with pytest.raises(ValueError, match="different Context"):
+            tv1.reset(token)
+        assert tv1.get_in(list(nursery.child_tasks)[0]) == 10
+        nursery.cancel_scope.cancel()
+
+
+def test_treevar_outside_run() -> None:
+    async def run_sync(fn: Any, *args: Any) -> Any:
+        return fn(*args)
+
+    tv1 = TreeVar("tv1", default=10)
+    for operation in (
+        tv1.get,
+        partial(tv1.get, 20),
+        partial(tv1.set, 30),
+        lambda: tv1.reset(trio.run(run_sync, tv1.set, 10)),
+        tv1.being(40).__enter__,
+    ):
+        with pytest.raises(RuntimeError, match="must be called from async context"):
+            operation()