pydocket 0.5.2__tar.gz → 0.6.1__tar.gz

{pydocket-0.5.2 → pydocket-0.6.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pydocket
-Version: 0.5.2
+Version: 0.6.1
 Summary: A distributed background task system for Python functions
 Project-URL: Homepage, https://github.com/chrisguidry/docket
 Project-URL: Bug Tracker, https://github.com/chrisguidry/docket/issues

{pydocket-0.5.2 → pydocket-0.6.1}/chaos/tasks.py

@@ -2,17 +2,27 @@ import logging
 import sys
 import time
 
-from docket import CurrentDocket, Docket, Retry, TaskKey
+from docket import CurrentDocket, Depends, Docket, Retry, TaskKey
 
 logger = logging.getLogger(__name__)
 
 
+async def greeting() -> str:
+    return "Hello, world"
+
+
+async def emphatic_greeting(greeting: str = Depends(greeting)) -> str:
+    return greeting + "!"
+
+
 async def hello(
+    greeting: str = Depends(emphatic_greeting),
     key: str = TaskKey(),
     docket: Docket = CurrentDocket(),
     retry: Retry = Retry(attempts=sys.maxsize),
 ):
     logger.info("Starting task %s", key)
+    logger.info("Greeting: %s", greeting)
     async with docket.redis() as redis:
         await redis.zadd("hello:received", {key: time.time()})
     logger.info("Finished task %s", key)

{pydocket-0.5.2 → pydocket-0.6.1}/examples/find_and_flood.py

@@ -16,7 +16,7 @@ async def find(
     perpetual: Perpetual = Perpetual(every=timedelta(seconds=3), automatic=True),
 ) -> None:
     for i in range(1, 10 + 1):
-        await docket.add(flood, key=str(i))(i)
+        await docket.add(flood, key=f"item-{i}")(i)
 
 
 async def flood(

{pydocket-0.5.2 → pydocket-0.6.1}/pyproject.toml

@@ -43,7 +43,7 @@ dev = [
     "opentelemetry-instrumentation-redis>=0.51b0",
     "opentelemetry-sdk>=1.30.0",
     "pre-commit>=4.1.0",
-    "pyright>=1.1.396",
+    "pyright>=1.1.398",
     "pytest>=8.3.4",
     "pytest-aio>=1.9.0",
     "pytest-cov>=6.0.0",

{pydocket-0.5.2 → pydocket-0.6.1}/src/docket/__init__.py

@@ -13,11 +13,13 @@ from .dependencies import (
     CurrentDocket,
     CurrentExecution,
     CurrentWorker,
+    Depends,
     ExponentialRetry,
     Perpetual,
     Retry,
     TaskKey,
     TaskLogger,
+    Timeout,
 )
 from .docket import Docket
 from .execution import Execution
@@ -36,5 +38,7 @@ __all__ = [
     "ExponentialRetry",
     "Logged",
     "Perpetual",
+    "Timeout",
+    "Depends",
     "__version__",
 ]

{pydocket-0.5.2 → pydocket-0.6.1}/src/docket/annotations.py

@@ -4,8 +4,14 @@ from typing import Any, Iterable, Mapping, Self
 
 
 class Annotation(abc.ABC):
+    _cache: dict[tuple[type[Self], inspect.Signature], Mapping[str, Self]] = {}
+
     @classmethod
     def annotated_parameters(cls, signature: inspect.Signature) -> Mapping[str, Self]:
+        key = (cls, signature)
+        if key in cls._cache:
+            return cls._cache[key]
+
         annotated: dict[str, Self] = {}
 
         for param_name, param in signature.parameters.items():
@@ -23,8 +29,25 @@ class Annotation(abc.ABC):
                 elif isinstance(arg_type, type) and issubclass(arg_type, cls):
                     annotated[param_name] = arg_type()
 
+        cls._cache[key] = annotated
         return annotated
 
 
 class Logged(Annotation):
     """Instructs docket to include arguments to this parameter in the log."""
+
+    length_only: bool = False
+
+    def __init__(self, length_only: bool = False) -> None:
+        self.length_only = length_only
+
+    def format(self, argument: Any) -> str:
+        if self.length_only:
+            if isinstance(argument, (dict, set)):
+                return f"{{len {len(argument)}}}"
+            elif isinstance(argument, tuple):
+                return f"(len {len(argument)})"
+            elif hasattr(argument, "__len__"):
+                return f"[len {len(argument)}]"
+
+        return repr(argument)

pydocket-0.6.1/src/docket/dependencies.py

@@ -0,0 +1,366 @@
+import abc
+import logging
+import time
+from contextlib import AsyncExitStack, asynccontextmanager
+from contextvars import ContextVar
+from datetime import timedelta
+from types import TracebackType
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    AsyncContextManager,
+    AsyncGenerator,
+    Awaitable,
+    Callable,
+    Counter,
+    Generic,
+    TypeVar,
+    cast,
+)
+
+from .docket import Docket
+from .execution import Execution, TaskFunction, get_signature
+
+if TYPE_CHECKING:  # pragma: no cover
+    from .worker import Worker
+
+
+class Dependency(abc.ABC):
+    single: bool = False
+
+    docket: ContextVar[Docket] = ContextVar("docket")
+    worker: ContextVar["Worker"] = ContextVar("worker")
+    execution: ContextVar[Execution] = ContextVar("execution")
+
+    @abc.abstractmethod
+    async def __aenter__(self) -> Any: ...  # pragma: no cover
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> bool: ...  # pragma: no cover
+
+
+class _CurrentWorker(Dependency):
+    async def __aenter__(self) -> "Worker":
+        return self.worker.get()
+
+
+def CurrentWorker() -> "Worker":
+    return cast("Worker", _CurrentWorker())
+
+
+class _CurrentDocket(Dependency):
+    async def __aenter__(self) -> Docket:
+        return self.docket.get()
+
+
+def CurrentDocket() -> Docket:
+    return cast(Docket, _CurrentDocket())
+
+
+class _CurrentExecution(Dependency):
+    async def __aenter__(self) -> Execution:
+        return self.execution.get()
+
+
+def CurrentExecution() -> Execution:
+    return cast(Execution, _CurrentExecution())
+
+
+class _TaskKey(Dependency):
+    async def __aenter__(self) -> str:
+        return self.execution.get().key
+
+
+def TaskKey() -> str:
+    return cast(str, _TaskKey())
+
+
+class _TaskLogger(Dependency):
+    async def __aenter__(self) -> logging.LoggerAdapter[logging.Logger]:
+        execution = self.execution.get()
+        logger = logging.getLogger(f"docket.task.{execution.function.__name__}")
+        return logging.LoggerAdapter(
+            logger,
+            {
+                **self.docket.get().labels(),
+                **self.worker.get().labels(),
+                **execution.specific_labels(),
+            },
+        )
+
+
+def TaskLogger() -> logging.LoggerAdapter[logging.Logger]:
+    return cast(logging.LoggerAdapter[logging.Logger], _TaskLogger())
+
+
+class Retry(Dependency):
+    single: bool = True
+
+    def __init__(
+        self, attempts: int | None = 1, delay: timedelta = timedelta(0)
+    ) -> None:
+        self.attempts = attempts
+        self.delay = delay
+        self.attempt = 1
+
+    async def __aenter__(self) -> "Retry":
+        execution = self.execution.get()
+        retry = Retry(attempts=self.attempts, delay=self.delay)
+        retry.attempt = execution.attempt
+        return retry
+
+
+class ExponentialRetry(Retry):
+    attempts: int
+
+    def __init__(
+        self,
+        attempts: int = 1,
+        minimum_delay: timedelta = timedelta(seconds=1),
+        maximum_delay: timedelta = timedelta(seconds=64),
+    ) -> None:
+        super().__init__(attempts=attempts, delay=minimum_delay)
+        self.minimum_delay = minimum_delay
+        self.maximum_delay = maximum_delay
+
+    async def __aenter__(self) -> "ExponentialRetry":
+        execution = self.execution.get()
+
+        retry = ExponentialRetry(
+            attempts=self.attempts,
+            minimum_delay=self.minimum_delay,
+            maximum_delay=self.maximum_delay,
+        )
+        retry.attempt = execution.attempt
+
+        if execution.attempt > 1:
+            backoff_factor = 2 ** (execution.attempt - 1)
+            calculated_delay = self.minimum_delay * backoff_factor
+
+            if calculated_delay > self.maximum_delay:
+                retry.delay = self.maximum_delay
+            else:
+                retry.delay = calculated_delay
+
+        return retry
+
+
+class Perpetual(Dependency):
+    single = True
+
+    every: timedelta
+    automatic: bool
+
+    args: tuple[Any, ...]
+    kwargs: dict[str, Any]
+
+    cancelled: bool
+
+    def __init__(
+        self,
+        every: timedelta = timedelta(0),
+        automatic: bool = False,
+    ) -> None:
+        """Declare a task that should be run perpetually.
+
+        Args:
+            every: The target interval between task executions.
+            automatic: If set, this task will be automatically scheduled during worker
+                startup and continually through the worker's lifespan.  This ensures
+                that the task will always be scheduled despite crashes and other
+                adverse conditions.  Automatic tasks must not require any arguments.
+        """
+        self.every = every
+        self.automatic = automatic
+        self.cancelled = False
+
+    async def __aenter__(self) -> "Perpetual":
+        execution = self.execution.get()
+        perpetual = Perpetual(every=self.every)
+        perpetual.args = execution.args
+        perpetual.kwargs = execution.kwargs
+        return perpetual
+
+    def cancel(self) -> None:
+        self.cancelled = True
+
+    def perpetuate(self, *args: Any, **kwargs: Any) -> None:
+        self.args = args
+        self.kwargs = kwargs
+
+
+class Timeout(Dependency):
+    single = True
+
+    base: timedelta
+
+    _deadline: float
+
+    def __init__(self, base: timedelta) -> None:
+        self.base = base
+
+    async def __aenter__(self) -> "Timeout":
+        timeout = Timeout(base=self.base)
+        timeout.start()
+        return timeout
+
+    def start(self) -> None:
+        self._deadline = time.monotonic() + self.base.total_seconds()
+
+    def expired(self) -> bool:
+        return time.monotonic() >= self._deadline
+
+    def remaining(self) -> timedelta:
+        return timedelta(seconds=self._deadline - time.monotonic())
+
+    def extend(self, by: timedelta | None = None) -> None:
+        if by is None:
+            by = self.base
+        self._deadline += by.total_seconds()
+
+
+R = TypeVar("R")
+
+DependencyFunction = Callable[..., Awaitable[R] | AsyncContextManager[R]]
+
+
+_parameter_cache: dict[
+    TaskFunction | DependencyFunction[Any],
+    dict[str, Dependency],
+] = {}
+
+
+def get_dependency_parameters(
+    function: TaskFunction | DependencyFunction[Any],
+) -> dict[str, Dependency]:
+    if function in _parameter_cache:
+        return _parameter_cache[function]
+
+    dependencies: dict[str, Dependency] = {}
+
+    signature = get_signature(function)
+
+    for parameter, param in signature.parameters.items():
+        if not isinstance(param.default, Dependency):
+            continue
+
+        dependencies[parameter] = param.default
+
+    _parameter_cache[function] = dependencies
+    return dependencies
+
+
+class _Depends(Dependency, Generic[R]):
+    dependency: DependencyFunction[R]
+
+    cache: ContextVar[dict[DependencyFunction[Any], Any]] = ContextVar("cache")
+    stack: ContextVar[AsyncExitStack] = ContextVar("stack")
+
+    def __init__(
+        self, dependency: Callable[[], Awaitable[R] | AsyncContextManager[R]]
+    ) -> None:
+        self.dependency = dependency
+
+    async def _resolve_parameters(
+        self,
+        function: TaskFunction | DependencyFunction[Any],
+    ) -> dict[str, Any]:
+        stack = self.stack.get()
+
+        arguments: dict[str, Any] = {}
+        parameters = get_dependency_parameters(function)
+
+        for parameter, dependency in parameters.items():
+            arguments[parameter] = await stack.enter_async_context(dependency)
+
+        return arguments
+
+    async def __aenter__(self) -> R:
+        cache = self.cache.get()
+
+        if self.dependency in cache:
+            return cache[self.dependency]
+
+        stack = self.stack.get()
+        arguments = await self._resolve_parameters(self.dependency)
+
+        value = self.dependency(**arguments)
+
+        if isinstance(value, AsyncContextManager):
+            value = await stack.enter_async_context(value)
+        else:
+            value = await value
+
+        cache[self.dependency] = value
+        return value
+
+
+def Depends(dependency: DependencyFunction[R]) -> R:
+    return cast(R, _Depends(dependency))
+
+
+D = TypeVar("D", bound=Dependency)
+
+
+def get_single_dependency_parameter_of_type(
+    function: TaskFunction, dependency_type: type[D]
+) -> D | None:
+    assert dependency_type.single, "Dependency must be single"
+    for _, dependency in get_dependency_parameters(function).items():
+        if isinstance(dependency, dependency_type):
+            return dependency
+    return None
+
+
+def get_single_dependency_of_type(
+    dependencies: dict[str, Dependency], dependency_type: type[D]
+) -> D | None:
+    assert dependency_type.single, "Dependency must be single"
+    for _, dependency in dependencies.items():
+        if isinstance(dependency, dependency_type):
+            return dependency
+    return None
+
+
+def validate_dependencies(function: TaskFunction) -> None:
+    parameters = get_dependency_parameters(function)
+
+    counts = Counter(type(dependency) for dependency in parameters.values())
+
+    for dependency_type, count in counts.items():
+        if dependency_type.single and count > 1:
+            raise ValueError(
+                f"Only one {dependency_type.__name__} dependency is allowed per task"
+            )
+
+
+@asynccontextmanager
+async def resolved_dependencies(
+    worker: "Worker", execution: Execution
+) -> AsyncGenerator[dict[str, Any], None]:
+    # Set context variables once at the beginning
+    Dependency.docket.set(worker.docket)
+    Dependency.worker.set(worker)
+    Dependency.execution.set(execution)
+
+    _Depends.cache.set({})
+
+    async with AsyncExitStack() as stack:
+        _Depends.stack.set(stack)
+
+        arguments: dict[str, Any] = {}
+
+        parameters = get_dependency_parameters(execution.function)
+        for parameter, dependency in parameters.items():
+            kwargs = execution.kwargs
+            if parameter in kwargs:
+                arguments[parameter] = kwargs[parameter]
+                continue
+
+            arguments[parameter] = await stack.enter_async_context(dependency)
+
+        yield arguments

{pydocket-0.5.2 → pydocket-0.6.1}/src/docket/docket.py

@@ -38,6 +38,7 @@ from .execution import (
     Strike,
     StrikeInstruction,
     StrikeList,
+    TaskFunction,
 )
 from .instrumentation import (
     REDIS_DISRUPTIONS,
@@ -57,7 +58,7 @@ tracer: trace.Tracer = trace.get_tracer(__name__)
 P = ParamSpec("P")
 R = TypeVar("R")
 
-TaskCollection = Iterable[Callable[..., Awaitable[Any]]]
+TaskCollection = Iterable[TaskFunction]
 
 RedisStreamID = bytes
 RedisMessageID = bytes
@@ -91,7 +92,7 @@ class RunningExecution(Execution):
         worker: str,
         started: datetime,
     ) -> None:
-        self.function: Callable[..., Awaitable[Any]] = execution.function
+        self.function: TaskFunction = execution.function
         self.args: tuple[Any, ...] = execution.args
         self.kwargs: dict[str, Any] = execution.kwargs
         self.when: datetime = execution.when
@@ -111,7 +112,7 @@ class DocketSnapshot:
 
 
 class Docket:
-    tasks: dict[str, Callable[..., Awaitable[Any]]]
+    tasks: dict[str, TaskFunction]
     strike_list: StrikeList
 
     _monitor_strikes_task: asyncio.Task[None]
@@ -197,7 +198,7 @@ class Docket:
         finally:
             await asyncio.shield(r.__aexit__(None, None, None))
 
-    def register(self, function: Callable[..., Awaitable[Any]]) -> None:
+    def register(self, function: TaskFunction) -> None:
         from .dependencies import validate_dependencies
 
         validate_dependencies(function)

{pydocket-0.5.2 → pydocket-0.6.1}/src/docket/execution.py

@@ -7,23 +7,40 @@ from typing import Any, Awaitable, Callable, Hashable, Literal, Mapping, Self, c
 
 import cloudpickle  # type: ignore[import]
 
+from opentelemetry import trace, propagate
+import opentelemetry.context
 
 from .annotations import Logged
+from docket.instrumentation import message_getter
 
 logger: logging.Logger = logging.getLogger(__name__)
 
+TaskFunction = Callable[..., Awaitable[Any]]
 Message = dict[bytes, bytes]
 
 
+_signature_cache: dict[Callable[..., Any], inspect.Signature] = {}
+
+
+def get_signature(function: Callable[..., Any]) -> inspect.Signature:
+    if function in _signature_cache:
+        return _signature_cache[function]
+
+    signature = inspect.signature(function)
+    _signature_cache[function] = signature
+    return signature
+
+
 class Execution:
     def __init__(
         self,
-        function: Callable[..., Awaitable[Any]],
+        function: TaskFunction,
         args: tuple[Any, ...],
         kwargs: dict[str, Any],
         when: datetime,
         key: str,
         attempt: int,
+        trace_context: opentelemetry.context.Context | None = None,
     ) -> None:
         self.function = function
         self.args = args
@@ -31,6 +48,7 @@ class Execution:
         self.when = when
         self.key = key
         self.attempt = attempt
+        self.trace_context = trace_context
 
     def as_message(self) -> Message:
         return {
@@ -43,9 +61,7 @@ class Execution:
         }
 
     @classmethod
-    def from_message(
-        cls, function: Callable[..., Awaitable[Any]], message: Message
-    ) -> Self:
+    def from_message(cls, function: TaskFunction, message: Message) -> Self:
         return cls(
             function=function,
             args=cloudpickle.loads(message[b"args"]),
@@ -53,6 +69,7 @@ class Execution:
             when=datetime.fromisoformat(message[b"when"].decode()),
             key=message[b"key"].decode(),
             attempt=int(message[b"attempt"].decode()),
+            trace_context=propagate.extract(message, getter=message_getter),
         )
 
     def general_labels(self) -> Mapping[str, str]:
@@ -68,28 +85,32 @@ class Execution:
 
     def call_repr(self) -> str:
         arguments: list[str] = []
-        signature = inspect.signature(self.function)
         function_name = self.function.__name__
 
+        signature = get_signature(self.function)
         logged_parameters = Logged.annotated_parameters(signature)
-
         parameter_names = list(signature.parameters.keys())
 
         for i, argument in enumerate(self.args[: len(parameter_names)]):
             parameter_name = parameter_names[i]
-            if parameter_name in logged_parameters:
-                arguments.append(repr(argument))
+            if logged := logged_parameters.get(parameter_name):
+                arguments.append(logged.format(argument))
             else:
                 arguments.append("...")
 
         for parameter_name, argument in self.kwargs.items():
-            if parameter_name in logged_parameters:
-                arguments.append(f"{parameter_name}={repr(argument)}")
+            if logged := logged_parameters.get(parameter_name):
+                arguments.append(f"{parameter_name}={logged.format(argument)}")
             else:
                 arguments.append(f"{parameter_name}=...")
 
         return f"{function_name}({', '.join(arguments)}){{{self.key}}}"
 
+    def incoming_span_links(self) -> list[trace.Link]:
+        initiating_span = trace.get_current_span(self.trace_context)
+        initiating_context = initiating_span.get_span_context()
+        return [trace.Link(initiating_context)] if initiating_context.is_valid else []
+
 
 class Operator(enum.StrEnum):
     EQUAL = "=="
@@ -217,10 +238,10 @@ class StrikeList:
         if function_name in self.task_strikes and not task_strikes:
             return True
 
-        sig = inspect.signature(execution.function)
+        signature = get_signature(execution.function)
 
         try:
-            bound_args = sig.bind(*execution.args, **execution.kwargs)
+            bound_args = signature.bind(*execution.args, **execution.kwargs)
             bound_args.apply_defaults()
         except TypeError:
             # If we can't make sense of the arguments, just assume the task is fine
@@ -265,6 +286,8 @@ class StrikeList:
                 case "between":  # pragma: no branch
                     lower, upper = strike_value
                     return lower <= value <= upper
+                case _:  # pragma: no cover
+                    raise ValueError(f"Unknown operator: {operator}")
         except (ValueError, TypeError):
             # If we can't make the comparison due to incompatible types, just log the
             # error and assume the task is not stricken