diff --git a/docs/usage.rst b/docs/usage.rst index cd91689..8296104 100644 --- a/docs/usage.rst +++ b/docs/usage.rst @@ -398,6 +398,42 @@ Programmatic registration: app.schedule.add_schedule(at(time(9, 0)) & on(Mon, Wed, Fri), "myapp.tasks:report") +Job mappings +~~~~~~~~~~~~ + +Jobs could be registered with a mapping, instead of being registered one-by-one. + +So, instead of: + +.. code-block:: python + async def generate_notifications( + kind: Literal["regular", "unique"], + ) -> None: + # ... + + + @broker.sched(at(time(18))) + async def generate_regular_notifications() -> None: + await generate_notifications(kind="regular") + + + @broker.sched(on_day(Thu) & at(time(18, 30))) + async def generate_unique_notifications() -> None: + await generate_notifications(kind="unique") + +One could write: + +.. code-block:: python + @broker.sched({ + at(time(18)): Payload(kwargs={"kind":"regular"}), + on_day(Thu) & at(time(18, 30)): Payload(kwargs={"kind":"unique"}), + }) + async def generate_notifications( + kind: Literal["regular", "unique"], + ) -> None: + # ... +``` + Serializers ----------- diff --git a/kuu/app.py b/kuu/app.py index 160ed21..24e2c75 100644 --- a/kuu/app.py +++ b/kuu/app.py @@ -2,7 +2,7 @@ import inspect from datetime import datetime, timedelta -from typing import Any, overload +from typing import Any, Mapping, overload from kuu._import import object_fqn from kuu._types import _FnAny, _FnAsync, _Wrap @@ -147,7 +147,7 @@ def wrap(fn: _FnAny[P, R]) -> Task[P, R]: def sched[**P, R]( self, - sched: Schedule, + sched: Schedule | Mapping[Schedule, Payload], args: Payload = Payload(), *, sched_id: str | None = None, @@ -157,7 +157,7 @@ def sched[**P, R]( ) -> _Wrap[P, R]: """Register function as a task with `Schedule` - :param sched: schedule to run task with + :param sched: schedule to run task with, or mapping of schedules to their corresponding arguments :param args: args for the task to be executed with :param sched_id: optional, specific sched_id for scheduler :param queue: destination queue; defaults to `Kuu.default_queue` @@ -168,15 +168,29 @@ def sched[**P, R]( def wrap(fn: _FnAny[P, R]) -> Task[P, R]: if not isinstance(fn, Task): fn = self.task(fn) - self.schedule.add_schedule( - sched, - fn, - args, - id=sched_id, - queue=queue, - headers=headers, - max_attempts=max_attempts, - ) + + if isinstance(sched, Schedule): + self.schedule.add_schedule( + sched, + fn, + args, + id=sched_id, + queue=queue, + headers=headers, + max_attempts=max_attempts, + ) + else: + for sched_, payload in sched.items(): + self.schedule.add_schedule( + sched_, + fn, + payload, + id=sched_id, + queue=queue, + headers=headers, + max_attempts=max_attempts, + ) + return fn return wrap diff --git a/tests/test_scheduler.py b/tests/test_scheduler.py index 663228e..b3f630e 100644 --- a/tests/test_scheduler.py +++ b/tests/test_scheduler.py @@ -1,6 +1,7 @@ from __future__ import annotations from datetime import datetime, time, timedelta, timezone +from typing import Literal import pytest @@ -208,6 +209,30 @@ async def morning_report(): assert "morning_report" in job.task_name +def test_kuu_sched_decorator_registers_mapping(): + app = Kuu(broker=MemoryBroker()) + + _mon_sched = at(time(9, 0)) & on(Mon) + _wed_sched = at(time(9, 0)) & on(Wed) + _fri_sched = at(time(9, 0)) & on(Fri) + + @app.sched( + { + _mon_sched: "mon", + _wed_sched: "wed", + _fri_sched: "fri", + } + ) + async def morning_report(day: Literal["mon", "wed", "fri"]) -> None: + pass + + assert len(app.schedule.jobs) == 3 + job = app.schedule.jobs[0] + assert isinstance(job, ScheduleJob) + assert "morning_report" in job.task_name + assert job.schedule == _mon_sched + + def test_kuu_every_decorator_accepts_optional_kwargs(): app = Kuu(broker=MemoryBroker())