dycw-utilities 0.132.4__tar.gz → 0.133.0__tar.gz

{dycw_utilities-0.132.4 → dycw_utilities-0.133.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dycw-utilities
-Version: 0.132.4
+Version: 0.133.0
 Author-email: Derek Wan <d.wan@icloud.com>
 License-File: LICENSE
 Requires-Python: >=3.12

{dycw_utilities-0.132.4 → dycw_utilities-0.133.0}/pyproject.toml

@@ -14,6 +14,7 @@ altair-test = [
   "img2pdf",
   "vl-convert-python",
 ]
+arq = ["arq >= 0.26.3, < 0.27"]
 atools = ["atools >= 0.14.2, < 0.15"]
 cachetools = ["cachetools >= 5.5.2, < 5.6"]
 click = ["click >= 8.2.1, < 8.3"]
@@ -74,7 +75,7 @@ pytest-regressions = ["pytest-regressions >= 2.8.0, < 2.9"]
 pytest-regressions-test = ["orjson", "polars-lts-cpu"]
 pytest-test = ["orjson", "pytest-rng", "pytest-rerunfailures"]
 python-dotenv = ["python-dotenv >= 1.1.0, < 1.2"]
-redis = ["redis >= 6.2.0, < 6.3", "orjson"]
+redis = ["redis >= 5.3.0, < 6.0", "orjson"]
 redis-test = ["pytest-rerunfailures"]
 reprlib-test = ["rich"]
 scipy = ["scipy >= 1.15.3, < 1.16"]
@@ -102,7 +103,7 @@ dependencies = [
 name = "dycw-utilities"
 readme = "README.md"
 requires-python = ">= 3.12"
-version = "0.132.4"
+version = "0.133.0"
 
 [project.optional-dependencies]
 logging = [
@@ -128,7 +129,7 @@ test = [
 # bump-my-version
 [tool.bumpversion]
 allow_dirty = true
-current_version = "0.132.4"
+current_version = "0.133.0"
 
 [[tool.bumpversion.files]]
 filename = "src/utilities/__init__.py"

dycw_utilities-0.133.0/src/tests/test_arq.py

@@ -0,0 +1,47 @@
+from __future__ import annotations
+
+from asyncio import sleep
+from typing import TYPE_CHECKING, Any, cast
+
+from hypothesis import given
+from hypothesis.strategies import integers
+
+from utilities.arq import Worker, cron_raw
+from utilities.iterables import one
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+    from arq.cron import CronJob
+    from arq.typing import WorkerCoroutine
+
+    from utilities.types import CallableCoroutine1
+
+
+class TestWorker:
+    @given(x=integers(), y=integers())
+    async def test_main(self, *, x: int, y: int) -> None:
+        async def func(x: int, y: int, /) -> int:
+            await sleep(0.01)
+            return x + y
+
+        class Example(Worker):
+            functions_raw: Sequence[CallableCoroutine1[Any]] = [func]
+
+        func_use = cast("WorkerCoroutine", one(Example.functions))
+        result = await func_use({}, x, y)
+        assert result == (x + y)
+
+    @given(x=integers(), y=integers())
+    async def test_cron(self, *, x: int, y: int) -> None:
+        async def func(x: int, y: int, /) -> int:
+            await sleep(0.01)
+            return x + y
+
+        class Example(Worker):
+            cron_jobs: Sequence[CronJob] | None = [cron_raw(func, args=(x, y))]
+
+        assert Example.cron_jobs is not None
+        cron_job = one(Example.cron_jobs)
+        result = await cron_job.coroutine({})
+        assert result == (x + y)

{dycw_utilities-0.132.4 → dycw_utilities-0.133.0}/src/utilities/__init__.py

@@ -1,3 +1,3 @@
 from __future__ import annotations
 
-__version__ = "0.132.4"
+__version__ = "0.133.0"

dycw_utilities-0.133.0/src/utilities/arq.py

@@ -0,0 +1,161 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from functools import wraps
+from itertools import chain
+from typing import TYPE_CHECKING, Any, ParamSpec, TypeVar, cast, override
+
+from arq.constants import default_queue_name, expires_extra_ms
+from arq.cron import cron
+
+if TYPE_CHECKING:
+    from collections.abc import Callable, Iterable, Sequence
+    from datetime import timezone
+
+    from arq.connections import ArqRedis, RedisSettings
+    from arq.cron import CronJob
+    from arq.jobs import Deserializer, Serializer
+    from arq.typing import (
+        OptionType,
+        SecondsTimedelta,
+        StartupShutdown,
+        WeekdayOptionType,
+        WorkerCoroutine,
+    )
+    from arq.worker import Function
+
+    from utilities.types import CallableCoroutine1, Coroutine1, StrMapping
+
+_P = ParamSpec("_P")
+_T = TypeVar("_T")
+
+
+##
+
+
+def cron_raw(
+    coroutine: CallableCoroutine1[Any],
+    /,
+    *,
+    name: str | None = None,
+    month: OptionType = None,
+    day: OptionType = None,
+    weekday: WeekdayOptionType = None,
+    hour: OptionType = None,
+    minute: OptionType = None,
+    second: OptionType = 0,
+    microsecond: int = 123_456,
+    run_at_startup: bool = False,
+    unique: bool = True,
+    job_id: str | None = None,
+    timeout: SecondsTimedelta | None = None,
+    keep_result: float | None = 0,
+    keep_result_forever: bool | None = False,
+    max_tries: int | None = 1,
+    args: Iterable[Any] | None = None,
+    kwargs: StrMapping | None = None,
+) -> CronJob:
+    """Create a cron job with a raw coroutine function."""
+    lifted = _lift_cron(
+        coroutine, *(() if args is None else args), **({} if kwargs is None else kwargs)
+    )
+    return cron(
+        lifted,
+        name=name,
+        month=month,
+        day=day,
+        weekday=weekday,
+        hour=hour,
+        minute=minute,
+        second=second,
+        microsecond=microsecond,
+        run_at_startup=run_at_startup,
+        unique=unique,
+        job_id=job_id,
+        timeout=timeout,
+        keep_result=keep_result,
+        keep_result_forever=keep_result_forever,
+        max_tries=max_tries,
+    )
+
+
+def _lift_cron(
+    func: Callable[_P, Coroutine1[_T]], *args: _P.args, **kwargs: _P.kwargs
+) -> WorkerCoroutine:
+    """Lift a coroutine function & call arg/kwargs for `cron`."""
+
+    @wraps(func)
+    async def wrapped(ctx: StrMapping, /) -> _T:
+        _ = ctx
+        return await func(*args, **kwargs)
+
+    return cast("Any", wrapped)
+
+
+##
+
+
+class _WorkerMeta(type):
+    @override
+    def __new__(
+        mcs: type[_WorkerMeta],
+        name: str,
+        bases: tuple[type, ...],
+        namespace: dict[str, Any],
+        /,
+    ) -> type[Worker]:
+        cls = cast("type[Worker]", super().__new__(mcs, name, bases, namespace))
+        cls.functions = tuple(chain(cls.functions, map(cls._lift, cls.functions_raw)))
+        return cls
+
+    @classmethod
+    def _lift(cls, func: Callable[_P, Coroutine1[_T]]) -> WorkerCoroutine:
+        """Lift a coroutine function to accept the required `ctx` argument."""
+
+        @wraps(func)
+        async def wrapped(ctx: StrMapping, *args: _P.args, **kwargs: _P.kwargs) -> _T:
+            _ = ctx
+            return await func(*args, **kwargs)
+
+        return cast("Any", wrapped)
+
+
+@dataclass(kw_only=True)
+class Worker(metaclass=_WorkerMeta):
+    """Base class for all workers."""
+
+    functions: Sequence[Function | WorkerCoroutine] = ()
+    functions_raw: Sequence[CallableCoroutine1[Any]] = ()
+    queue_name: str | None = default_queue_name
+    cron_jobs: Sequence[CronJob] | None = None
+    redis_settings: RedisSettings | None = None
+    redis_pool: ArqRedis | None = None
+    burst: bool = False
+    on_startup: StartupShutdown | None = None
+    on_shutdown: StartupShutdown | None = None
+    on_job_start: StartupShutdown | None = None
+    on_job_end: StartupShutdown | None = None
+    after_job_end: StartupShutdown | None = None
+    handle_signals: bool = True
+    job_completion_wait: int = 0
+    max_jobs: int = 10
+    job_timeout: SecondsTimedelta = 300
+    keep_result: SecondsTimedelta = 3600
+    keep_result_forever: bool = False
+    poll_delay: SecondsTimedelta = 0.5
+    queue_read_limit: int | None = None
+    max_tries: int = 5
+    health_check_interval: SecondsTimedelta = 3600
+    health_check_key: str | None = None
+    ctx: dict[Any, Any] | None = None
+    retry_jobs: bool = True
+    allow_abort_jobs: bool = False
+    max_burst_jobs: int = -1
+    job_serializer: Serializer | None = None
+    job_deserializer: Deserializer | None = None
+    expires_extra_ms: int = expires_extra_ms
+    timezone: timezone | None = None
+    log_results: bool = True
+
+
+__all__ = ["Worker", "cron"]

{dycw_utilities-0.132.4 → dycw_utilities-0.133.0}/src/utilities/types.py

@@ -72,12 +72,16 @@ type Coroutine1[_T] = Coroutine[Any, Any, _T]
 type MaybeAwaitable[_T] = _T | Awaitable[_T]
 type MaybeCallableEvent = MaybeCallable[Event]
 type MaybeCoroutine1[_T] = _T | Coroutine1[_T]
+type CallableCoroutine1[_T] = Callable[..., Coroutine1[_T]]
 
 
 # callable
 TCallable = TypeVar("TCallable", bound=Callable[..., Any])
 TCallable1 = TypeVar("TCallable1", bound=Callable[..., Any])
 TCallable2 = TypeVar("TCallable2", bound=Callable[..., Any])
+TCallableCoroutine1 = TypeVar(
+    "TCallableCoroutine1", bound=Callable[..., Coroutine1[Any]]
+)
 TCallableMaybeCoroutine1None = TypeVar(
     "TCallableMaybeCoroutine1None", bound=Callable[..., MaybeCoroutine1[None]]
 )
@@ -292,6 +296,7 @@ type TimeZoneLike = (
 
 
 __all__ = [
+    "CallableCoroutine1",
     "Coroutine1",
     "Dataclass",
     "DateDeltaLike",
@@ -346,6 +351,7 @@ __all__ = [
     "TCallable",
     "TCallable1",
     "TCallable2",
+    "TCallableCoroutine1",
     "TCallableMaybeCoroutine1None",
     "TDataclass",
     "TEnum",