Split cancellable_sleep

In response to feedback, I've added `raise_if_cancelled()` to replace `cancellable_sleep(None)`.

In response to the same feedback, I now handle the error if no invocation ID is available, so we simply perform a regular time.sleep.

I've also deleted a couple of defunct print statements.

Author: rwb27

src/labthings_fastapi/__init__.py

@@ -35,6 +35,7 @@
 from .invocation_contexts import (
     get_invocation_logger,
     cancellable_sleep,
+    raise_if_cancelled,
     ThreadWithInvocationID,
 )
 
@@ -62,5 +63,6 @@
     "ThingClient",
     "get_invocation_logger",
     "cancellable_sleep",
+    "raise_if_cancelled",
     "ThreadWithInvocationID",
 ]

src/labthings_fastapi/invocation_contexts.py

@@ -19,6 +19,7 @@
 from contextlib import contextmanager
 import logging
 from threading import Event, Thread
+import time
 from typing import Any, Callable
 from typing_extensions import Self
 from uuid import UUID, uuid4
@@ -191,8 +192,8 @@ def get_cancel_event(id: UUID | None = None) -> CancelEvent:
     return CancelEvent.get_for_id(id)
 
 
-def cancellable_sleep(interval: float | None) -> None:
-    """Sleep for a specified time, allowing cancellation.
+def cancellable_sleep(interval: float) -> None:
+    r"""Sleep for a specified time, allowing cancellation.
 
     This function should be called from action functions instead of
     `time.sleep` to allow them to be cancelled. Usually, this
@@ -206,26 +207,38 @@ def cancellable_sleep(interval: float | None) -> None:
         This function uses `.Event.wait` internally, which suffers
         from timing errors on some platforms: it may have error of
         around 10-20ms. If that's a problem, consider using
-        `time.sleep` instead. ``lt.cancellable_sleep(None)`` may then
+        `time.sleep` instead. ``lt.raise_if_cancelled()`` may then
         be used to allow cancellation.
 
-    This function may only be called from an action thread, as it
-    depends on the invocation ID being available from a context variable.
-    Use `.set_invocation_id` to make it available outside of an action
-    thread.
+    If this function is called from outside of an action thread, it
+    will revert to `time.sleep`\ .
+
+    :param interval: The length of time to wait for, in seconds.
+    """
+    try:
+        event = get_cancel_event()
+        event.sleep(interval)
+    except NoInvocationContextError:
+        time.sleep(interval)
 
-    If ``interval`` is set to None, we do not call `.Event.wait` but
-    instead we simply check whether the event is set.
 
-    :param interval: The length of time to sleep for, in seconds. If it
-        is `None` we won't wait, but we will still check for a cancel
-        event, and raise the exception if it is set.
+def raise_if_cancelled() -> None:
+    """Raise an exception if the current invocation has been cancelled.
+
+    This function checks for cancellation events and, if the current
+    action invocation has been cancelled, it will raise an
+    `.InvocationCancelledError` to signal the thread to terminate.
+    It is equivalent to `.cancellable_sleep` but without waiting any
+    time.
+
+    If called outside of an invocation context, this function does
+    nothing, and will not raise an error.
     """
-    event = get_cancel_event()
-    if interval is None:
+    try:
+        event = get_cancel_event()
         event.raise_if_set()
-    else:
-        event.sleep(interval)
+    except NoInvocationContextError:
+        pass
 
 
 def get_invocation_logger(id: UUID | None = None) -> logging.Logger:
@@ -335,11 +348,9 @@ def join_and_propagate_cancel(self, poll_interval: float = 0.2) -> None:
         cancellation: InvocationCancelledError | None = None
         self._polls = 0
         self._attempted_cancels = 0
-        print(f"Checking for cancellation of invocation {get_invocation_id()}")
-        print(f"so we can cancel {self.invocation_id}")
         while self.is_alive():
             try:
-                cancellable_sleep(None)
+                raise_if_cancelled()
                 self._polls += 1
             except InvocationCancelledError as e:
                 # Propagate the cancellation signal to the thread

tests/test_invocation_contexts.py

@@ -134,10 +134,10 @@ def test_cancel_event():
 
 def test_cancellable_sleep():
     """Check the module-level cancellable sleep."""
-    with pytest.raises(exc.NoInvocationContextError):
-        ic.cancellable_sleep(1)
-    with pytest.raises(exc.NoInvocationContextError):
-        ic.cancellable_sleep(None)
+    # with no invocation context, the function should wait
+    # and there should be no error.
+    with assert_takes_time(0.02, 0.04):
+        ic.cancellable_sleep(0.02)
 
     with ic.fake_invocation_context():
         event = ic.get_cancel_event()
@@ -146,20 +146,30 @@ def test_cancellable_sleep():
         with assert_takes_time(0.02, 0.04):
             ic.cancellable_sleep(0.02)
 
-        # passing `None` should return immediately.
-        with assert_takes_time(None, 0.002):
-            ic.cancellable_sleep(None)
-
         # check an exception gets raised and reset if appropriate
         event.set()
         with pytest.raises(exc.InvocationCancelledError):
             ic.cancellable_sleep(1)
         assert not event.is_set()
 
+
+def test_raise_if_cancelled():
+    """Check the module-level cancellable sleep."""
+    # the function should return immediately.
+    with assert_takes_time(None, 0.002):
+        ic.raise_if_cancelled()
+
+    with ic.fake_invocation_context():
+        event = ic.get_cancel_event()
+
+        # the function should return immediately.
+        with assert_takes_time(None, 0.002):
+            ic.raise_if_cancelled()
+
         # check an exception gets raised and reset if appropriate
         event.set()
         with pytest.raises(exc.InvocationCancelledError):
-            ic.cancellable_sleep(None)
+            ic.raise_if_cancelled()
         assert not event.is_set()
 
 
@@ -179,19 +189,6 @@ def test_invocation_logger():
         assert logger.name.endswith(str(id))
 
 
-def run_function_in_thread_and_propagate_cancellation(func, *args):
-    """Run a function in a ThreadWithInvocationID."""
-    t = ic.ThreadWithInvocationID(target=func, args=args)
-    t.start()
-    try:
-        t.join_and_propagate_cancel(0.005)
-    except exc.InvocationCancelledError:
-        # We still want to return the finished thread if it's
-        # cancelled.
-        pass
-    return t
-
-
 def test_thread_with_invocation_id():
     """Test our custom thread subclass makes a new ID and can be cancelled."""
     ids = []
@@ -204,6 +201,9 @@ def test_thread_with_invocation_id():
     assert t.exception is None
     assert t.result is None
 
+
+def test_thread_with_invocation_id_cancel():
+    """Test the custom thread subclass responds to cancellation."""
     # Check cancellable sleep works in the thread
     t = ic.ThreadWithInvocationID(target=ic.cancellable_sleep, args=[1])
     assert isinstance(t.invocation_id, uuid.UUID)
@@ -213,20 +213,43 @@ def test_thread_with_invocation_id():
         t.join()
     assert isinstance(t.exception, exc.InvocationCancelledError)
 
-    # Check we capture the return value
+
+def test_thread_with_invocation_id_return_value():
+    """Check we capture the return value when running in a ThreadWithInvocationID."""
     t = ic.ThreadWithInvocationID(target=lambda: True)
     t.start()
     t.join()
     assert t.exception is None
     assert t.result is True
 
+
+def run_function_in_thread_and_propagate_cancellation(func, *args):
+    """Run a function in a ThreadWithInvocationID."""
+    t = ic.ThreadWithInvocationID(target=func, args=args)
+    t.start()
+    try:
+        t.join_and_propagate_cancel(1)
+    except exc.InvocationCancelledError:
+        # We still want to return the finished thread if it's
+        # cancelled.
+        pass
+    return t
+
+
+def test_thread_with_invocation_id_cancellation_propagates():
+    """Check that a cancel event can propagate to our thread.
+
+    ``join_and_propagate_cancellation`` should cancel the spawned thread if
+    the parent thread is cancelled while it's waiting for the spawned thread
+    to join.
+    """
     # Check we can propagate cancellation.
     # First, we run `cancellable_sleep` and check it doesn't cancel
     with ic.fake_invocation_context():
         # First test our function - there is only one thread here, and we
         # check it finishes and doesn't error.
         t = run_function_in_thread_and_propagate_cancellation(
-            ic.cancellable_sleep, 0.001
+            ic.cancellable_sleep, 0.02
         )
         assert isinstance(t, ic.ThreadWithInvocationID)
         assert not t.is_alive()