Add in_ and at helpers to allow for directly requesting a retry

Author: bunchesofdonald

src/docket/dependencies.py

@@ -3,7 +3,7 @@
 import time
 from contextlib import AsyncExitStack, asynccontextmanager
 from contextvars import ContextVar
-from datetime import timedelta
+from datetime import datetime, timedelta
 from types import TracebackType
 from typing import (
     TYPE_CHECKING,
@@ -14,11 +14,13 @@
     Callable,
     Counter,
     Generic,
+    NoReturn,
     TypeVar,
     cast,
 )
 
 from .docket import Docket
+from .exceptions import ForcedRetry
 from .execution import Execution, TaskFunction, get_signature
 
 if TYPE_CHECKING:  # pragma: no cover
@@ -222,6 +224,17 @@ async def __aenter__(self) -> "Retry":
         retry.attempt = execution.attempt
         return retry
 
+    def at(self, when: datetime) -> NoReturn:
+        now = datetime.now()
+        diff = when - now
+        diff = diff if diff.total_seconds() >= 0 else timedelta(0)
+
+        self.in_(diff)
+
+    def in_(self, when: timedelta) -> NoReturn:
+        self.delay = when
+        raise ForcedRetry()
+
 
 class ExponentialRetry(Retry):
     """Configures exponential retries for a task.  You can specify the total number

src/docket/exceptions.py

@@ -0,0 +1,2 @@
+class ForcedRetry(Exception):
+    """Raised when a task requests a retry via `in_` or `at`"""

tests/test_dependencies.py

@@ -1,4 +1,5 @@
 import logging
+from datetime import datetime, timedelta
 
 import pytest
 
@@ -95,6 +96,125 @@ async def the_task(
     assert calls == 2
 
 
+async def test_user_can_request_a_retry_in_timedelta_time(
+    docket: Docket, worker: Worker
+):
+    calls = 0
+    first_call_time = None
+    second_call_time = None
+
+    async def the_task(
+        a: str,
+        b: str,
+        retry: Retry = Retry(attempts=2),
+    ):
+        assert a == "a"
+        assert b == "b"
+
+        nonlocal calls
+        calls += 1
+
+        nonlocal first_call_time
+        if not first_call_time:
+            first_call_time = datetime.now()
+            retry.in_(timedelta(seconds=0.5))
+        else:
+            nonlocal second_call_time
+            second_call_time = datetime.now()
+
+    await docket.add(the_task)("a", "b")
+
+    await worker.run_until_finished()
+
+    assert calls == 2
+
+    assert isinstance(first_call_time, datetime)
+    assert isinstance(second_call_time, datetime)
+
+    delay = second_call_time - first_call_time
+    assert delay.total_seconds() > 0 < 1
+
+
+async def test_user_can_request_a_retry_at_a_specific_time(
+    docket: Docket, worker: Worker
+):
+    calls = 0
+    first_call_time = None
+    second_call_time = None
+
+    async def the_task(
+        a: str,
+        b: str,
+        retry: Retry = Retry(attempts=2),
+    ):
+        assert a == "a"
+        assert b == "b"
+
+        nonlocal calls
+        calls += 1
+
+        nonlocal first_call_time
+        if not first_call_time:
+            when = datetime.now() + timedelta(seconds=0.5)
+            first_call_time = datetime.now()
+            retry.at(when)
+        else:
+            nonlocal second_call_time
+            second_call_time = datetime.now()
+
+    await docket.add(the_task)("a", "b")
+
+    await worker.run_until_finished()
+
+    assert calls == 2
+
+    assert isinstance(first_call_time, datetime)
+    assert isinstance(second_call_time, datetime)
+
+    delay = second_call_time - first_call_time
+    assert delay.total_seconds() > 0 < 1
+
+
+async def test_user_can_request_a_retry_at_a_specific_time_in_the_past(
+    docket: Docket, worker: Worker
+):
+    calls = 0
+    first_call_time = None
+    second_call_time = None
+
+    async def the_task(
+        a: str,
+        b: str,
+        retry: Retry = Retry(attempts=2),
+    ):
+        assert a == "a"
+        assert b == "b"
+
+        nonlocal calls
+        calls += 1
+
+        nonlocal first_call_time
+        if not first_call_time:
+            when = datetime.now() - timedelta(days=1)
+            first_call_time = datetime.now()
+            retry.at(when)
+        else:
+            nonlocal second_call_time
+            second_call_time = datetime.now()
+
+    await docket.add(the_task)("a", "b")
+
+    await worker.run_until_finished()
+
+    assert calls == 2
+
+    assert isinstance(first_call_time, datetime)
+    assert isinstance(second_call_time, datetime)
+
+    delay = second_call_time - first_call_time
+    assert delay.total_seconds() > 0 < 1
+
+
 async def test_dependencies_error_for_missing_task_argument(
     docket: Docket, worker: Worker, caplog: pytest.LogCaptureFixture
 ):