pydocket 0.5.1__tar.gz → 0.6.0__tar.gz

{pydocket-0.5.1 → pydocket-0.6.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pydocket
-Version: 0.5.1
+Version: 0.6.0
 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.1 → pydocket-0.6.0}/chaos/driver.py

@@ -23,12 +23,6 @@ from .tasks import toxic
 
 logging.getLogger().setLevel(logging.INFO)
 
-# Quiets down the testcontainers logger
-testcontainers_logger = logging.getLogger("testcontainers.core.container")
-testcontainers_logger.setLevel(logging.ERROR)
-testcontainers_logger = logging.getLogger("testcontainers.core.waiting_utils")
-testcontainers_logger.setLevel(logging.ERROR)
-
 console = logging.StreamHandler(stream=sys.stdout)
 console.setFormatter(
     logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")

{pydocket-0.5.1 → pydocket-0.6.0}/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.6.0/examples/common.py

@@ -0,0 +1,69 @@
+import asyncio
+import os
+import socket
+from contextlib import asynccontextmanager
+from typing import AsyncGenerator
+
+from docker import DockerClient
+
+
+@asynccontextmanager
+async def run_redis(version: str) -> AsyncGenerator[str, None]:
+    def get_free_port() -> int:
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+            s.bind(("", 0))
+            return s.getsockname()[1]
+
+    port = get_free_port()
+
+    client = DockerClient.from_env()
+    container = client.containers.run(
+        f"redis:{version}",
+        detach=True,
+        ports={"6379/tcp": port},
+        auto_remove=True,
+    )
+
+    # Wait for Redis to be ready
+    for line in container.logs(stream=True):
+        if b"Ready to accept connections" in line:
+            break
+
+    url = f"redis://localhost:{port}/0"
+    print("***** Redis is running on %s *****", url)
+    try:
+        yield url
+    finally:
+        container.stop()
+
+
+async def run_example_workers(workers: int, concurrency: int, tasks: str):
+    async with run_redis("7.4.2") as redis_url:
+        processes = [
+            await asyncio.create_subprocess_exec(
+                "docket",
+                "worker",
+                "--name",
+                f"worker-{i}",
+                "--url",
+                redis_url,
+                "--tasks",
+                tasks,
+                "--concurrency",
+                str(concurrency),
+                env={
+                    **os.environ,
+                    "PYTHONPATH": os.path.abspath(
+                        os.path.join(os.path.dirname(__file__), "..")
+                    ),
+                },
+            )
+            for i in range(workers)
+        ]
+        try:
+            await asyncio.gather(*[p.wait() for p in processes])
+        except asyncio.CancelledError:
+            for p in processes:
+                p.kill()
+        finally:
+            await asyncio.gather(*[p.wait() for p in processes])

pydocket-0.6.0/examples/find_and_flood.py

@@ -0,0 +1,39 @@
+import asyncio
+from datetime import timedelta
+from logging import Logger, LoggerAdapter
+from typing import Annotated
+
+from docket import Docket
+from docket.annotations import Logged
+from docket.dependencies import CurrentDocket, Perpetual, TaskLogger
+
+from .common import run_example_workers
+
+
+async def find(
+    docket: Docket = CurrentDocket(),
+    logger: LoggerAdapter[Logger] = TaskLogger(),
+    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)
+
+
+async def flood(
+    item: Annotated[int, Logged],
+    logger: LoggerAdapter[Logger] = TaskLogger(),
+) -> None:
+    logger.info("Working on %s", item)
+
+
+tasks = [find, flood]
+
+
+if __name__ == "__main__":
+    asyncio.run(
+        run_example_workers(
+            workers=3,
+            concurrency=8,
+            tasks="examples.find_and_flood:tasks",
+        )
+    )

{pydocket-0.5.1 → pydocket-0.6.0}/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.1 → pydocket-0.6.0}/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.1 → pydocket-0.6.0}/src/docket/annotations.py

@@ -28,3 +28,19 @@ class Annotation(abc.ABC):
 
 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.0/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