pydocket 0.0.1__py3-none-any.whl

docket/__init__.py

@@ -0,0 +1,24 @@
+"""
+docket - A distributed background task system for Python functions.
+
+docket focuses on scheduling future work as seamlessly and efficiently as immediate work.
+"""
+
+from importlib.metadata import version
+
+__version__ = version("pydocket")
+
+from .dependencies import CurrentDocket, CurrentWorker, Retry
+from .docket import Docket
+from .execution import Execution
+from .worker import Worker
+
+__all__ = [
+    "Docket",
+    "Worker",
+    "Execution",
+    "CurrentDocket",
+    "CurrentWorker",
+    "Retry",
+    "__version__",
+]

docket/__main__.py

@@ -0,0 +1,3 @@
+from docket.cli import app
+
+app()

docket/cli.py

@@ -0,0 +1,23 @@
+import typer
+
+from docket import __version__
+
+app: typer.Typer = typer.Typer(
+    help="Docket - A distributed background task system for Python functions",
+    add_completion=True,
+    no_args_is_help=True,
+)
+
+
+@app.command(
+    help="Start a worker to process tasks",
+)
+def worker() -> None:
+    print("TODO: start the worker")
+
+
+@app.command(
+    help="Print the version of Docket",
+)
+def version() -> None:
+    print(__version__)

docket/dependencies.py

@@ -0,0 +1,77 @@
+import abc
+import inspect
+from datetime import timedelta
+from typing import Any, Awaitable, Callable, Counter, cast
+
+from .docket import Docket
+from .execution import Execution
+from .worker import Worker
+
+
+class Dependency(abc.ABC):
+    single: bool = False
+
+    @abc.abstractmethod
+    def __call__(
+        self, docket: Docket, worker: Worker, execution: Execution
+    ) -> Any: ...  # pragma: no cover
+
+
+class _CurrentWorker(Dependency):
+    def __call__(self, docket: Docket, worker: Worker, execution: Execution) -> Worker:
+        return worker
+
+
+def CurrentWorker() -> Worker:
+    return cast(Worker, _CurrentWorker())
+
+
+class _CurrentDocket(Dependency):
+    def __call__(self, docket: Docket, worker: Worker, execution: Execution) -> Docket:
+        return docket
+
+
+def CurrentDocket() -> Docket:
+    return cast(Docket, _CurrentDocket())
+
+
+class Retry(Dependency):
+    single: bool = True
+
+    def __init__(self, attempts: int = 1, delay: timedelta = timedelta(0)) -> None:
+        self.attempts = attempts
+        self.delay = delay
+        self.attempt = 1
+
+    def __call__(self, docket: Docket, worker: Worker, execution: Execution) -> "Retry":
+        retry = Retry(attempts=self.attempts, delay=self.delay)
+        retry.attempt = execution.attempt
+        return retry
+
+
+def get_dependency_parameters(
+    function: Callable[..., Awaitable[Any]],
+) -> dict[str, Dependency]:
+    dependencies: dict[str, Any] = {}
+
+    signature = inspect.signature(function)
+
+    for param_name, param in signature.parameters.items():
+        if not isinstance(param.default, Dependency):
+            continue
+
+        dependencies[param_name] = param.default
+
+    return dependencies
+
+
+def validate_dependencies(function: Callable[..., Awaitable[Any]]) -> 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"
+            )

docket/docket.py

@@ -0,0 +1,178 @@
+from contextlib import asynccontextmanager
+from datetime import datetime, timezone
+from types import TracebackType
+from typing import (
+    Any,
+    AsyncGenerator,
+    Awaitable,
+    Callable,
+    ParamSpec,
+    Self,
+    TypeVar,
+    overload,
+)
+from uuid import uuid4
+
+from redis.asyncio import Redis
+
+from .execution import Execution
+
+P = ParamSpec("P")
+R = TypeVar("R")
+
+
+class Docket:
+    tasks: dict[str, Callable[..., Awaitable[Any]]]
+
+    def __init__(
+        self,
+        name: str = "docket",
+        host: str = "localhost",
+        port: int = 6379,
+        db: int = 0,
+        password: str | None = None,
+    ) -> None:
+        self.name = name
+        self.host = host
+        self.port = port
+        self.db = db
+        self.password = password
+
+    async def __aenter__(self) -> Self:
+        self.tasks = {}
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> None:
+        pass
+
+    @asynccontextmanager
+    async def redis(self) -> AsyncGenerator[Redis, None]:
+        async with Redis(
+            host=self.host,
+            port=self.port,
+            db=self.db,
+            password=self.password,
+            single_connection_client=True,
+        ) as redis:
+            yield redis
+
+    def register(self, function: Callable[..., Awaitable[Any]]) -> None:
+        from .dependencies import validate_dependencies
+
+        validate_dependencies(function)
+
+        self.tasks[function.__name__] = function
+
+    @overload
+    def add(
+        self,
+        function: Callable[P, Awaitable[R]],
+        when: datetime | None = None,
+        key: str | None = None,
+    ) -> Callable[P, Awaitable[Execution]]: ...  # pragma: no cover
+
+    @overload
+    def add(
+        self,
+        function: str,
+        when: datetime | None = None,
+        key: str | None = None,
+    ) -> Callable[..., Awaitable[Execution]]: ...  # pragma: no cover
+
+    def add(
+        self,
+        function: Callable[P, Awaitable[R]] | str,
+        when: datetime | None = None,
+        key: str | None = None,
+    ) -> Callable[..., Awaitable[Execution]]:
+        if isinstance(function, str):
+            function = self.tasks[function]
+        else:
+            self.register(function)
+
+        if when is None:
+            when = datetime.now(timezone.utc)
+
+        if key is None:
+            key = f"{function.__name__}:{uuid4()}"
+
+        async def scheduler(*args: P.args, **kwargs: P.kwargs) -> Execution:
+            execution = Execution(function, args, kwargs, when, key, attempt=1)
+            await self.schedule(execution)
+            return execution
+
+        return scheduler
+
+    @overload
+    def replace(
+        self,
+        function: Callable[P, Awaitable[R]],
+        when: datetime,
+        key: str,
+    ) -> Callable[P, Awaitable[Execution]]: ...  # pragma: no cover
+
+    @overload
+    def replace(
+        self,
+        function: str,
+        when: datetime,
+        key: str,
+    ) -> Callable[..., Awaitable[Execution]]: ...  # pragma: no cover
+
+    def replace(
+        self,
+        function: Callable[P, Awaitable[R]] | str,
+        when: datetime,
+        key: str,
+    ) -> Callable[..., Awaitable[Execution]]:
+        if isinstance(function, str):
+            function = self.tasks[function]
+
+        async def scheduler(*args: P.args, **kwargs: P.kwargs) -> Execution:
+            execution = Execution(function, args, kwargs, when, key, attempt=1)
+            await self.cancel(key)
+            await self.schedule(execution)
+            return execution
+
+        return scheduler
+
+    @property
+    def queue_key(self) -> str:
+        return f"{self.name}:queue"
+
+    @property
+    def stream_key(self) -> str:
+        return f"{self.name}:stream"
+
+    def parked_task_key(self, key: str) -> str:
+        return f"{self.name}:{key}"
+
+    async def schedule(self, execution: Execution) -> None:
+        message: dict[bytes, bytes] = execution.as_message()
+        key = execution.key
+        when = execution.when
+
+        async with self.redis() as redis:
+            # if the task is already in the queue, retain it
+            if await redis.zscore(self.queue_key, key) is not None:
+                return
+
+            if when <= datetime.now(timezone.utc):
+                await redis.xadd(self.stream_key, message)
+            else:
+                async with redis.pipeline() as pipe:
+                    pipe.hset(self.parked_task_key(key), mapping=message)
+                    pipe.zadd(self.queue_key, {key: when.timestamp()})
+                    await pipe.execute()
+
+    async def cancel(self, key: str) -> None:
+        async with self.redis() as redis:
+            async with redis.pipeline() as pipe:
+                pipe.delete(self.parked_task_key(key))
+                pipe.zrem(self.queue_key, key)
+                await pipe.execute()

docket/execution.py

@@ -0,0 +1,47 @@
+from datetime import datetime
+from typing import Any, Awaitable, Callable, Self
+
+import cloudpickle
+
+Message = dict[bytes, bytes]
+
+
+class Execution:
+    def __init__(
+        self,
+        function: Callable[..., Awaitable[Any]],
+        args: tuple[Any, ...],
+        kwargs: dict[str, Any],
+        when: datetime,
+        key: str,
+        attempt: int,
+    ) -> None:
+        self.function = function
+        self.args = args
+        self.kwargs = kwargs
+        self.when = when
+        self.key = key
+        self.attempt = attempt
+
+    def as_message(self) -> Message:
+        return {
+            b"key": self.key.encode(),
+            b"when": self.when.isoformat().encode(),
+            b"function": self.function.__name__.encode(),
+            b"args": cloudpickle.dumps(self.args),
+            b"kwargs": cloudpickle.dumps(self.kwargs),
+            b"attempt": str(self.attempt).encode(),
+        }
+
+    @classmethod
+    def from_message(
+        cls, function: Callable[..., Awaitable[Any]], message: Message
+    ) -> Self:
+        return cls(
+            function=function,
+            args=cloudpickle.loads(message[b"args"]),
+            kwargs=cloudpickle.loads(message[b"kwargs"]),
+            when=datetime.fromisoformat(message[b"when"].decode()),
+            key=message[b"key"].decode(),
+            attempt=int(message[b"attempt"].decode()),
+        )

docket/worker.py

@@ -0,0 +1,244 @@
+import logging
+import sys
+from datetime import datetime, timezone
+from types import TracebackType
+from typing import TYPE_CHECKING, Any, Protocol, Self, Sequence, TypeVar, cast
+from uuid import uuid4
+
+from redis import RedisError
+
+from .docket import Docket, Execution
+
+logger: logging.Logger = logging.getLogger(__name__)
+
+RedisStreamID = bytes
+RedisMessageID = bytes
+RedisMessage = dict[bytes, bytes]
+RedisStream = tuple[RedisStreamID, Sequence[tuple[RedisMessageID, RedisMessage]]]
+RedisReadGroupResponse = Sequence[RedisStream]
+
+if TYPE_CHECKING:  # pragma: no cover
+    from .dependencies import Dependency
+
+D = TypeVar("D", bound="Dependency")
+
+
+class _stream_due_tasks(Protocol):
+    async def __call__(
+        self, keys: list[str], args: list[str | float]
+    ) -> tuple[int, int]: ...  # pragma: no cover
+
+
+class Worker:
+    name: str
+    docket: Docket
+
+    prefetch_count: int = 10
+
+    def __init__(self, docket: Docket) -> None:
+        self.name = f"worker:{uuid4()}"
+        self.docket = docket
+
+    async def __aenter__(self) -> Self:
+        async with self.docket.redis() as redis:
+            try:
+                await redis.xgroup_create(
+                    groupname=self.consumer_group_name,
+                    name=self.docket.stream_key,
+                    id="0-0",
+                    mkstream=True,
+                )
+            except RedisError as e:
+                assert "BUSYGROUP" in repr(e)
+
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> None:
+        pass
+
+    @property
+    def consumer_group_name(self) -> str:
+        return "docket"
+
+    @property
+    def _log_context(self) -> dict[str, str]:
+        return {
+            "queue_key": self.docket.queue_key,
+            "stream_key": self.docket.stream_key,
+        }
+
+    async def run_until_current(self) -> None:
+        async with self.docket.redis() as redis:
+            stream_due_tasks: _stream_due_tasks = cast(
+                _stream_due_tasks,
+                redis.register_script(
+                    # Lua script to atomically move scheduled tasks to the stream
+                    # KEYS[1]: queue key (sorted set)
+                    # KEYS[2]: stream key
+                    # ARGV[1]: current timestamp
+                    # ARGV[2]: docket name prefix
+                    """
+                local total_work = redis.call('ZCARD', KEYS[1])
+                local due_work = 0
+                local tasks = redis.call('ZRANGEBYSCORE', KEYS[1], 0, ARGV[1])
+
+                for i, key in ipairs(tasks) do
+                    local hash_key = ARGV[2] .. ":" .. key
+                    local task_data = redis.call('HGETALL', hash_key)
+
+                    if #task_data > 0 then
+                        local task = {}
+                        for j = 1, #task_data, 2 do
+                            task[task_data[j]] = task_data[j+1]
+                        end
+
+                        redis.call('XADD', KEYS[2], '*',
+                            'key', task['key'],
+                            'when', task['when'],
+                            'function', task['function'],
+                            'args', task['args'],
+                            'kwargs', task['kwargs'],
+                            'attempt', task['attempt']
+                        )
+                        redis.call('DEL', hash_key)
+                        due_work = due_work + 1
+                    end
+                end
+
+                if due_work > 0 then
+                    redis.call('ZREMRANGEBYSCORE', KEYS[1], 0, ARGV[1])
+                end
+
+                return {total_work, due_work}
+                """
+                ),
+            )
+
+            total_work, due_work = sys.maxsize, 0
+            while total_work:
+                now = datetime.now(timezone.utc)
+                total_work, due_work = await stream_due_tasks(
+                    keys=[self.docket.queue_key, self.docket.stream_key],
+                    args=[now.timestamp(), self.docket.name],
+                )
+                logger.info(
+                    "Moved %d/%d due tasks from %s to %s",
+                    due_work,
+                    total_work,
+                    self.docket.queue_key,
+                    self.docket.stream_key,
+                    extra=self._log_context,
+                )
+
+                response: RedisReadGroupResponse = await redis.xreadgroup(
+                    groupname=self.consumer_group_name,
+                    consumername=self.name,
+                    streams={self.docket.stream_key: ">"},
+                    count=self.prefetch_count,
+                    block=10,
+                )
+                for _, messages in response:
+                    for message_id, message in messages:
+                        await self._execute(message)
+
+                        # When executing a task, there's always a chance that it was
+                        # either retried or it scheduled another task, so let's give
+                        # ourselves one more iteration of the loop to handle that.
+                        total_work += 1
+
+                        async with redis.pipeline() as pipe:
+                            pipe.xack(
+                                self.docket.stream_key,
+                                self.consumer_group_name,
+                                message_id,
+                            )
+                            pipe.xdel(
+                                self.docket.stream_key,
+                                message_id,
+                            )
+                            await pipe.execute()
+
+    async def _execute(self, message: RedisMessage) -> None:
+        execution = Execution.from_message(
+            self.docket.tasks[message[b"function"].decode()],
+            message,
+        )
+
+        logger.info(
+            "Executing task %s with args %s and kwargs %s",
+            execution.key,
+            execution.args,
+            execution.kwargs,
+            extra={
+                **self._log_context,
+                "function": execution.function.__name__,
+            },
+        )
+
+        dependencies = self._get_dependencies(execution)
+
+        try:
+            await execution.function(
+                *execution.args,
+                **{
+                    **execution.kwargs,
+                    **dependencies,
+                },
+            )
+        except Exception:
+            logger.exception(
+                "Error executing task %s",
+                execution.key,
+                extra=self._log_context,
+            )
+            await self._retry_if_requested(execution, dependencies)
+
+    def _get_dependencies(
+        self,
+        execution: Execution,
+    ) -> dict[str, Any]:
+        from .dependencies import get_dependency_parameters
+
+        parameters = get_dependency_parameters(execution.function)
+
+        dependencies: dict[str, Any] = {}
+
+        for param_name, dependency in parameters.items():
+            # If the argument is already provided, skip it, which allows users to call
+            # the function directly with the arguments they want.
+            if param_name in execution.kwargs:
+                dependencies[param_name] = execution.kwargs[param_name]
+                continue
+
+            dependencies[param_name] = dependency(self.docket, self, execution)
+
+        return dependencies
+
+    async def _retry_if_requested(
+        self,
+        execution: Execution,
+        dependencies: dict[str, Any],
+    ) -> None:
+        from .dependencies import Retry
+
+        retries = [retry for retry in dependencies.values() if isinstance(retry, Retry)]
+        if not retries:
+            return
+
+        retry = retries[0]
+
+        if execution.attempt < retry.attempts:
+            execution.when = datetime.now(timezone.utc) + retry.delay
+            execution.attempt += 1
+            await self.docket.schedule(execution)
+        else:
+            logger.error(
+                "Task %s failed after %d attempts",
+                execution.key,
+                retry.attempts,
+            )

pydocket-0.0.1.dist-info/METADATA

@@ -0,0 +1,31 @@
+Metadata-Version: 2.4
+Name: pydocket
+Version: 0.0.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
+Author-email: Chris Guidry <guid@omg.lol>
+License: # Released under MIT License
+        
+        Copyright (c) 2025 Chris Guidry.
+        
+        Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
+        
+        The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
+        
+        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+License-File: LICENSE
+Classifier: Development Status :: 4 - Beta
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Requires-Python: >=3.12
+Requires-Dist: cloudpickle>=3.1.1
+Requires-Dist: redis>=5.2.1
+Requires-Dist: typer>=0.15.1
+Description-Content-Type: text/markdown
+
+docket is a distributed background task system for Python functions with a focus
+on the scheduling of future work as seamlessly and efficiency as immediate work.

pydocket-0.0.1.dist-info/RECORD

@@ -0,0 +1,13 @@
+docket/__init__.py,sha256=0gzwWxJcDqU7nEZ4rHrWYk7WqEgW4Mz77E7rbjpyCHw,526
+docket/__main__.py,sha256=Vkuh7aJ-Bl7QVpVbbkUksAd_hn05FiLmWbc-8kbhZQ4,34
+docket/cli.py,sha256=ty8CirvLDvvOuOs7MHW0SYMjmbQq9rqUgPCq-HqW_6g,434
+docket/dependencies.py,sha256=yd_sIv3y69Czo_DyHh3aNtLJGFOMjg8jjJ701QN04-Q,2124
+docket/docket.py,sha256=sIgSvPUX4HG4EN_OxPSN7xCGu7DZkk5oZsuBVDp5XcU,4942
+docket/execution.py,sha256=3HT1GOeg76RMILiq06bpDZITHrqjVQ_j9FU6fuY4jqw,1375
+docket/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+docket/worker.py,sha256=N-nftCveJELKnICp1lgmzG5-r1QCYkGRxi8FcTjrLSc,8204
+pydocket-0.0.1.dist-info/METADATA,sha256=IYrB8lUbawu5I9Buh6uVJ68kJh47GUfph3F1wxxN78w,2084
+pydocket-0.0.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+pydocket-0.0.1.dist-info/entry_points.txt,sha256=4WOk1nUlBsUT5O3RyMci2ImuC5XFswuopElYcLHtD5k,47
+pydocket-0.0.1.dist-info/licenses/LICENSE,sha256=YuVWU_ZXO0K_k2FG8xWKe5RGxV24AhJKTvQmKfqXuyk,1087
+pydocket-0.0.1.dist-info/RECORD,,

pydocket-0.0.1.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.27.0
+Root-Is-Purelib: true
+Tag: py3-none-any

pydocket-0.0.1.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+docket = docket.__main__:app

pydocket-0.0.1.dist-info/licenses/LICENSE

@@ -0,0 +1,9 @@
+# Released under MIT License
+
+Copyright (c) 2025 Chris Guidry.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.