move_on_ and fail_ functions accepts shield kwarg

docs/source/reference-core.rst

@@ -431,8 +431,7 @@ attribute to :data:`True`::
        try:
            await conn.send_hello_msg()
        finally:
-           with trio.move_on_after(CLEANUP_TIMEOUT) as cleanup_scope:
-               cleanup_scope.shield = True
+           with trio.move_on_after(CLEANUP_TIMEOUT, shield=True) as cleanup_scope:
                await conn.send_goodbye_msg()
 
 So long as you're inside a scope with ``shield = True`` set, then

newsfragments/3049.feature.rst

@@ -0,0 +1 @@
+`trio.move_on_at`, `trio.move_on_after`, `trio.fail_at` and `trio.fail_after` now accepts shield as a kwarg.

src/trio/_tests/test_timeouts.py

@@ -75,6 +75,66 @@ async def sleep_3() -> None:
     await check_takes_about(sleep_3, TARGET)
 
 
+@slow
+async def test_move_on_after_shields_from_outer() -> None:
+    duration = 0.1
+
+    async def task() -> None:
+        with _core.CancelScope() as outer, move_on_after(TARGET, shield=True) as inner:
+            outer.cancel()
+            # The outer scope is cancelled, but this task is protected by the
+            # shield, so it manages to get to sleep for 'duration' seconds
+            await trio.sleep(duration)
+            # now when we unshield, it should abort the sleep.
+            inner.shield = False
+
+    # Check that the task takes only about 'duration' seconds
+    await check_takes_about(task, duration)
+
+
+@slow
+async def test_move_on_after_moves_on_even_if_shielded() -> None:
+    async def task() -> None:
+        with _core.CancelScope() as outer, move_on_after(TARGET, shield=True):
+            outer.cancel()
+            # The outer scope is cancelled, but this task is protected by the
+            # shield, so it manages to get to sleep until deadline is met
+            await sleep_forever()
+
+    await check_takes_about(task, TARGET)
+
+
+@slow
+async def test_fail_after_exits_shields_from_outer() -> None:
+    duration = 0.1
+
+    async def task() -> None:
+        with _core.CancelScope() as outer, fail_after(TARGET, shield=True) as inner:
+            outer.cancel()
+            # The outer scope is cancelled, but this task is protected by the
+            # shield, so it manages to get to sleep for 'duration' seconds
+            await trio.sleep(duration)
+            # now when we unshield, it should abort the sleep.
+            inner.shield = False
+
+    # Check that the task takes only about 'duration' seconds
+    await check_takes_about(task, duration)
+
+
+@slow
+async def test_fail_after_fails_even_if_shielded() -> None:
+    async def task() -> None:
+        with pytest.raises(TooSlowError), _core.CancelScope() as outer, fail_after(
+            TARGET, shield=True
+        ):
+            outer.cancel()
+            # The outer scope is cancelled, but this task is protected by the
+            # shield, so it manages to get to sleep until deadline is met
+            await sleep_forever()
+
+    await check_takes_about(task, TARGET)
+
+
 @slow
 async def test_fail() -> None:
     async def sleep_4() -> None:

src/trio/_timeouts.py

@@ -7,36 +7,38 @@
 import trio
 
 
-def move_on_at(deadline: float) -> trio.CancelScope:
+def move_on_at(deadline: float, *, shield: bool = False) -> trio.CancelScope:
     """Use as a context manager to create a cancel scope with the given
     absolute deadline.
 
     Args:
       deadline (float): The deadline.
+      shield (bool): Whether the cancel scope created is shielded.
 
     Raises:
       ValueError: if deadline is NaN.
 
     """
     if math.isnan(deadline):
         raise ValueError("deadline must not be NaN")
-    return trio.CancelScope(deadline=deadline)
+    return trio.CancelScope(deadline=deadline, shield=shield)
 
 
-def move_on_after(seconds: float) -> trio.CancelScope:
+def move_on_after(seconds: float, *, shield: bool = False) -> trio.CancelScope:
     """Use as a context manager to create a cancel scope whose deadline is
     set to now + *seconds*.
 
     Args:
       seconds (float): The timeout.
+      shield (bool): Whether the cancel scope created is shielded.
 
     Raises:
       ValueError: if timeout is less than zero or NaN.
 
     """
     if seconds < 0:
         raise ValueError("timeout must be non-negative")
-    return move_on_at(trio.current_time() + seconds)
+    return move_on_at(trio.current_time() + seconds, shield=shield)
 
 
 async def sleep_forever() -> None:
@@ -96,7 +98,7 @@ class TooSlowError(Exception):
 
 # workaround for PyCharm not being able to infer return type from @contextmanager
 # see https://youtrack.jetbrains.com/issue/PY-36444/PyCharm-doesnt-infer-types-when-using-contextlib.contextmanager-decorator
-def fail_at(deadline: float) -> AbstractContextManager[trio.CancelScope]:  # type: ignore[misc]
+def fail_at(deadline: float, *, shield: bool = False) -> AbstractContextManager[trio.CancelScope]:  # type: ignore[misc]
     """Creates a cancel scope with the given deadline, and raises an error if it
     is actually cancelled.
 
@@ -110,14 +112,15 @@ def fail_at(deadline: float) -> AbstractContextManager[trio.CancelScope]:  # typ
 
     Args:
       deadline (float): The deadline.
+      shield (bool): Whether the cancel scope created is shielded.
 
     Raises:
       TooSlowError: if a :exc:`Cancelled` exception is raised in this scope
         and caught by the context manager.
       ValueError: if deadline is NaN.
 
     """
-    with move_on_at(deadline) as scope:
+    with move_on_at(deadline, shield=shield) as scope:
         yield scope
     if scope.cancelled_caught:
         raise TooSlowError
@@ -127,7 +130,9 @@ def fail_at(deadline: float) -> AbstractContextManager[trio.CancelScope]:  # typ
     fail_at = contextmanager(fail_at)
 
 
-def fail_after(seconds: float) -> AbstractContextManager[trio.CancelScope]:
+def fail_after(
+    seconds: float, *, shield: bool = False
+) -> AbstractContextManager[trio.CancelScope]:
     """Creates a cancel scope with the given timeout, and raises an error if
     it is actually cancelled.
 
@@ -140,6 +145,7 @@ def fail_after(seconds: float) -> AbstractContextManager[trio.CancelScope]:
 
     Args:
       seconds (float): The timeout.
+      shield (bool): Whether the cancel scope created is shielded.
 
     Raises:
       TooSlowError: if a :exc:`Cancelled` exception is raised in this scope
@@ -149,4 +155,4 @@ def fail_after(seconds: float) -> AbstractContextManager[trio.CancelScope]:
     """
     if seconds < 0:
         raise ValueError("timeout must be non-negative")
-    return fail_at(trio.current_time() + seconds)
+    return fail_at(trio.current_time() + seconds, shield=shield)