From 9ee445e54e2eb3a2e8ef7bb302ee9c57961cc4c0 Mon Sep 17 00:00:00 2001 From: Lura Skye Revuwution Date: Sun, 1 Aug 2021 02:57:10 +0100 Subject: [PATCH 1/5] Add API for changing the context of a running task. --- trio/_core/_context.py | 25 +++++++++++++++++++++++++ trio/_core/tests/test_context.py | 31 +++++++++++++++++++++++++++++++ 2 files changed, 56 insertions(+) create mode 100644 trio/_core/_context.py create mode 100644 trio/_core/tests/test_context.py diff --git a/trio/_core/_context.py b/trio/_core/_context.py new file mode 100644 index 0000000000..c504fdaec0 --- /dev/null +++ b/trio/_core/_context.py @@ -0,0 +1,25 @@ +from contextvars import Context + +from async_generator import asynccontextmanager + +from . import _run + + +@asynccontextmanager +async def change_context(context: Context): + """Asynchronous context manager to change the :class:`contextvars.Context` + for the current task. + + """ + task = _run.current_task() + saved_context = task.context + task.context = context + + try: + await _run.checkpoint() + yield + finally: + task.context = saved_context + # I assume this is right? + await _run.cancel_shielded_checkpoint() + diff --git a/trio/_core/tests/test_context.py b/trio/_core/tests/test_context.py new file mode 100644 index 0000000000..706468e078 --- /dev/null +++ b/trio/_core/tests/test_context.py @@ -0,0 +1,31 @@ +from contextvars import ContextVar, copy_context + +from .._run import CancelScope +from .._context import change_context + +var = ContextVar("var") + + +async def test_context_change(): + var.set("abc") + + new_context = copy_context() + async with change_context(new_context): + assert var.get() == "abc" + var.set("def") + assert var.get() == "def" + + assert var.get() == "abc" + + +async def test_context_change_with_cancellation(): + var.set("qqq") + + new_context = copy_context() + with CancelScope() as scope: + async with change_context(new_context): + var.set("ghj") + assert var.get() == "ghj" + scope.cancel() + + assert var.get() == "qqq" From 506d7c8b6a75d0da3411191b74b12a2144b6efa3 Mon Sep 17 00:00:00 2001 From: Lura Skye Revuwution Date: Sun, 1 Aug 2021 03:08:28 +0100 Subject: [PATCH 2/5] Export change_context. --- trio/__init__.py | 1 + trio/_core/__init__.py | 2 ++ 2 files changed, 3 insertions(+) diff --git a/trio/__init__.py b/trio/__init__.py index a50ec33310..112fd9ec14 100644 --- a/trio/__init__.py +++ b/trio/__init__.py @@ -32,6 +32,7 @@ BrokenResourceError, EndOfChannel, Nursery, + change_context, ) from ._timeouts import ( diff --git a/trio/_core/__init__.py b/trio/_core/__init__.py index 2bd0c74e67..ddab0b19b8 100644 --- a/trio/_core/__init__.py +++ b/trio/_core/__init__.py @@ -75,6 +75,8 @@ from ._mock_clock import MockClock +from ._context import change_context + # Windows imports if sys.platform == "win32": from ._run import ( From 885b88820c876674596877fccb8a1ae98c8267ae Mon Sep 17 00:00:00 2001 From: Lura Skye Revuwution Date: Sun, 1 Aug 2021 03:08:56 +0100 Subject: [PATCH 3/5] Newsfragment + doc. --- docs/source/reference-core.rst | 6 ++++++ newsfragments/2086.feature.rst | 2 ++ trio/_core/_context.py | 2 +- 3 files changed, 9 insertions(+), 1 deletion(-) create mode 100644 newsfragments/2086.feature.rst diff --git a/docs/source/reference-core.rst b/docs/source/reference-core.rst index 514b5072c3..7b7a951da9 100644 --- a/docs/source/reference-core.rst +++ b/docs/source/reference-core.rst @@ -1020,6 +1020,12 @@ Example output (yours may differ slightly): request 0: Helper task b finished request 0: Request received finished +You can change the current :class:`contextvars.Context` a task is running +in with a helper context manager. + +.. autofunction:: change_context + :async-with: + For more information, read the `contextvar docs `__. diff --git a/newsfragments/2086.feature.rst b/newsfragments/2086.feature.rst new file mode 100644 index 0000000000..f0407f01f0 --- /dev/null +++ b/newsfragments/2086.feature.rst @@ -0,0 +1,2 @@ +Add an API to allow changing the current Context a task is running +in. \ No newline at end of file diff --git a/trio/_core/_context.py b/trio/_core/_context.py index c504fdaec0..bc571d0681 100644 --- a/trio/_core/_context.py +++ b/trio/_core/_context.py @@ -6,7 +6,7 @@ @asynccontextmanager -async def change_context(context: Context): +async def change_context(context): """Asynchronous context manager to change the :class:`contextvars.Context` for the current task. From 4f9736d340cd3e1b9150401023773180993fde2d Mon Sep 17 00:00:00 2001 From: Lura Skye Revuwution Date: Sun, 1 Aug 2021 03:10:36 +0100 Subject: [PATCH 4/5] Blackify. --- trio/_core/_context.py | 1 - 1 file changed, 1 deletion(-) diff --git a/trio/_core/_context.py b/trio/_core/_context.py index bc571d0681..d61003f9a9 100644 --- a/trio/_core/_context.py +++ b/trio/_core/_context.py @@ -22,4 +22,3 @@ async def change_context(context): task.context = saved_context # I assume this is right? await _run.cancel_shielded_checkpoint() - From 24b381456e57476cf4c120b8baad55647bef19eb Mon Sep 17 00:00:00 2001 From: Lura Skye Date: Thu, 23 Dec 2021 04:22:51 +0000 Subject: [PATCH 5/5] Address review comments. --- docs/source/reference-core.rst | 6 ------ docs/source/reference-lowlevel.rst | 14 ++++++++++++++ newsfragments/2086.feature.rst | 4 ++-- trio/__init__.py | 1 - trio/_core/__init__.py | 2 +- trio/_core/_context.py | 11 ++++------- trio/_core/tests/test_context.py | 6 +++--- trio/lowlevel.py | 1 + 8 files changed, 25 insertions(+), 20 deletions(-) diff --git a/docs/source/reference-core.rst b/docs/source/reference-core.rst index 7b7a951da9..514b5072c3 100644 --- a/docs/source/reference-core.rst +++ b/docs/source/reference-core.rst @@ -1020,12 +1020,6 @@ Example output (yours may differ slightly): request 0: Helper task b finished request 0: Request received finished -You can change the current :class:`contextvars.Context` a task is running -in with a helper context manager. - -.. autofunction:: change_context - :async-with: - For more information, read the `contextvar docs `__. diff --git a/docs/source/reference-lowlevel.rst b/docs/source/reference-lowlevel.rst index e0a604bebc..24ecaa199e 100644 --- a/docs/source/reference-lowlevel.rst +++ b/docs/source/reference-lowlevel.rst @@ -540,6 +540,20 @@ Task API :func:`wait_task_rescheduled` for details.) +Low-level context tools +======================= + +You can change the current :class:`contextvars.Context` a task is running +in with a helper context manager. + +.. autofunction:: set_current_context + :async-with: + +.. note:: + + This context manager introduces a scheduling point but not a cancellation + point. + .. _guest-mode: Using "guest mode" to run Trio on top of other event loops diff --git a/newsfragments/2086.feature.rst b/newsfragments/2086.feature.rst index f0407f01f0..4402996d29 100644 --- a/newsfragments/2086.feature.rst +++ b/newsfragments/2086.feature.rst @@ -1,2 +1,2 @@ -Add an API to allow changing the current Context a task is running -in. \ No newline at end of file +Add the :func:`trio.lowlevel.set_current_context` API to allow changing the +current :class:`contextvars.Context` a task is running in. \ No newline at end of file diff --git a/trio/__init__.py b/trio/__init__.py index 112fd9ec14..a50ec33310 100644 --- a/trio/__init__.py +++ b/trio/__init__.py @@ -32,7 +32,6 @@ BrokenResourceError, EndOfChannel, Nursery, - change_context, ) from ._timeouts import ( diff --git a/trio/_core/__init__.py b/trio/_core/__init__.py index ddab0b19b8..52fd2de0ab 100644 --- a/trio/_core/__init__.py +++ b/trio/_core/__init__.py @@ -75,7 +75,7 @@ from ._mock_clock import MockClock -from ._context import change_context +from ._context import set_current_context # Windows imports if sys.platform == "win32": diff --git a/trio/_core/_context.py b/trio/_core/_context.py index d61003f9a9..4821d1440e 100644 --- a/trio/_core/_context.py +++ b/trio/_core/_context.py @@ -1,14 +1,12 @@ -from contextvars import Context - from async_generator import asynccontextmanager from . import _run @asynccontextmanager -async def change_context(context): - """Asynchronous context manager to change the :class:`contextvars.Context` - for the current task. +async def set_current_context(context): + """Returns an asynchronous context manager to change the + :class:`contextvars.Context` for the current task. """ task = _run.current_task() @@ -16,9 +14,8 @@ async def change_context(context): task.context = context try: - await _run.checkpoint() + await _run.cancel_shielded_checkpoint() yield finally: task.context = saved_context - # I assume this is right? await _run.cancel_shielded_checkpoint() diff --git a/trio/_core/tests/test_context.py b/trio/_core/tests/test_context.py index 706468e078..b968cbc983 100644 --- a/trio/_core/tests/test_context.py +++ b/trio/_core/tests/test_context.py @@ -1,7 +1,7 @@ from contextvars import ContextVar, copy_context from .._run import CancelScope -from .._context import change_context +from .._context import set_current_context var = ContextVar("var") @@ -10,7 +10,7 @@ async def test_context_change(): var.set("abc") new_context = copy_context() - async with change_context(new_context): + async with set_current_context(new_context): assert var.get() == "abc" var.set("def") assert var.get() == "def" @@ -23,7 +23,7 @@ async def test_context_change_with_cancellation(): new_context = copy_context() with CancelScope() as scope: - async with change_context(new_context): + async with set_current_context(new_context): var.set("ghj") assert var.get() == "ghj" scope.cancel() diff --git a/trio/lowlevel.py b/trio/lowlevel.py index f30fdcc4bd..447e17ad27 100644 --- a/trio/lowlevel.py +++ b/trio/lowlevel.py @@ -44,6 +44,7 @@ notify_closing, start_thread_soon, start_guest_run, + set_current_context, ) from ._subprocess import open_process