diff --git a/procrastinate/blueprints.py b/procrastinate/blueprints.py index e39619189..9b072b112 100644 --- a/procrastinate/blueprints.py +++ b/procrastinate/blueprints.py @@ -4,6 +4,7 @@ import logging import sys from typing import TYPE_CHECKING, Callable, Literal, Union, cast, overload +from zoneinfo import ZoneInfo from typing_extensions import Concatenate, ParamSpec, TypeVar, Unpack @@ -195,6 +196,7 @@ def add_tasks_from(self, blueprint: Blueprint, *, namespace: str) -> None: task=periodic_task.task, cron=periodic_task.cron, periodic_id=periodic_task.periodic_id, + tzinfo=periodic_task.tzinfo, configure_kwargs=periodic_task.configure_kwargs, ) @@ -359,6 +361,7 @@ def periodic( *, cron: str, periodic_id: str = "", + tzinfo: str | None | ZoneInfo = None, **configure_kwargs: Unpack[ConfigureTaskOptions], ): """ @@ -371,11 +374,17 @@ def periodic( Cron-like string. Optionally add a 6th column for seconds. periodic_id : Task name suffix. Used to distinguish periodic tasks with different kwargs. - **kwargs : - Additional parameters are passed to `Task.configure`. + tzinfo : + Timezone in which the cron expression should be interpreted. Accepts a + timezone name string (e.g., "Africa/Blantyre"), a `zoneinfo.ZoneInfo` instance, + or `None`. When `None` (the default), the underlying `croniter` library + will interpret the schedule in UTC (the current behaviour). + **configure_kwargs : + Additional parameters are passed to ``Task.configure``. """ + return self.periodic_registry.periodic_decorator( - cron=cron, periodic_id=periodic_id, **configure_kwargs + cron=cron, periodic_id=periodic_id, tzinfo=tzinfo, **configure_kwargs ) def will_configure_task(self) -> None: diff --git a/procrastinate/periodic.py b/procrastinate/periodic.py index 94f4cd83a..d50debbf5 100644 --- a/procrastinate/periodic.py +++ b/procrastinate/periodic.py @@ -6,6 +6,7 @@ import time from collections.abc import Iterable from typing import Callable, Generic, cast +from zoneinfo import ZoneInfo, ZoneInfoNotFoundError import attr import croniter @@ -35,10 +36,26 @@ class PeriodicTask(Generic[P, R, Args]): cron: str periodic_id: str configure_kwargs: tasks.ConfigureTaskOptions + tzinfo: str | None | ZoneInfo = None @cached_property def croniter(self) -> croniter.croniter: - return croniter.croniter(self.cron) + croniter_instance = croniter.croniter(self.cron) + # croniter sets the timezone info object in + croniter_instance.tzinfo = self.get_tzinfo() + return croniter_instance + + def get_tzinfo(self): + tzinfo = self.tzinfo + if isinstance(tzinfo, str): + try: + return ZoneInfo(tzinfo) + except (ZoneInfoNotFoundError, ValueError): + logger.error(f"{tzinfo} is not a valid timezone.") + return None + if isinstance(tzinfo, ZoneInfo): + return tzinfo + return None TaskAtTime = tuple[PeriodicTask, int] @@ -52,6 +69,7 @@ def periodic_decorator( self, cron: str, periodic_id: str, + tzinfo: str | None | ZoneInfo = None, **configure_kwargs: Unpack[tasks.ConfigureTaskOptions], ) -> Callable[[tasks.Task[P, R, Concatenate[int, Args]]], tasks.Task[P, R, Args]]: """ @@ -68,6 +86,7 @@ def wrapper( cron=cron, periodic_id=periodic_id, configure_kwargs=configure_kwargs, + tzinfo=tzinfo, ) return cast(tasks.Task[P, R, Args], task) @@ -79,6 +98,7 @@ def register_task( cron: str, periodic_id: str, configure_kwargs: tasks.ConfigureTaskOptions, + tzinfo: str | None | ZoneInfo = None, ) -> PeriodicTask[P, R, Concatenate[int, Args]]: key = (task.name, periodic_id) if key in self.periodic_tasks: @@ -99,12 +119,12 @@ def register_task( "kwargs": str(configure_kwargs), }, ) - self.periodic_tasks[key] = periodic_task = PeriodicTask( task=task, cron=cron, periodic_id=periodic_id, configure_kwargs=configure_kwargs, + tzinfo=tzinfo, ) return periodic_task