Added timezone-aware schedule support. (#198)

Author: s3rius

poetry.lock

@@ -43,6 +43,7 @@ uvloop = { version = ">=0.16.0,<1", optional = true }
 # For hot-reload.
 watchdog = { version = "^2.1.9", optional = true }
 gitignore-parser = { version = "^0", optional = true }
+pytz = "*"
 
 [tool.poetry.dev-dependencies]
 pytest = "^7.1.2"
@@ -62,6 +63,8 @@ wemake-python-styleguide = "^0.18.0"
 tox = "^4.6.4"
 freezegun = "^1.2.2"
 pytest-mock = "^3.11.1"
+tzlocal = "^5.0.1"
+types-tzlocal = "^5.0.1.1"
 
 [tool.poetry.extras]
 zmq = ["pyzmq"]

pyproject.toml

@@ -3,6 +3,7 @@
 from logging import basicConfig, getLevelName, getLogger
 from typing import List
 
+import pytz
 from pycron import is_now
 
 from taskiq.cli.scheduler.args import SchedulerArgs
@@ -12,6 +13,22 @@
 logger = getLogger(__name__)
 
 
+def to_tz_aware(time: datetime) -> datetime:
+    """
+    Convert datetime to timezone aware.
+
+    This function takes a datetime and if
+    timezone was not yet specified, it will
+    be set to UTC.
+
+    :param time: time to make timezone aware.
+    :return: timezone aware time.
+    """
+    if time.tzinfo is None:
+        return time.replace(tzinfo=pytz.UTC)
+    return time
+
+
 async def schedules_updater(
     scheduler: TaskiqScheduler,
     current_schedules: List[ScheduledTask],
@@ -55,9 +72,18 @@ def should_run(task: ScheduledTask) -> bool:
     :return: True if task must be sent.
     """
     if task.cron is not None:
-        return is_now(task.cron, datetime.utcnow())
+        now = datetime.now(tz=pytz.UTC)
+        # If user specified cron offset we apply it.
+        # If it's timedelta, we simply add the delta to current time.
+        if task.cron_offset and isinstance(task.cron_offset, timedelta):
+            now += task.cron_offset
+        # If timezone was specified as string we convert it timzone
+        # offset and then apply.
+        elif task.cron_offset and isinstance(task.cron_offset, str):
+            now = now.astimezone(pytz.timezone(task.cron_offset))
+        return is_now(task.cron, now)
     if task.time is not None:
-        return task.time <= datetime.utcnow()
+        return to_tz_aware(task.time) <= datetime.now(tz=pytz.UTC)
     return False
 
 

taskiq/cli/scheduler/run.py

@@ -43,6 +43,7 @@ async def get_schedules(self) -> List["ScheduledTask"]:
                         kwargs=schedule.get("kwargs", {}),
                         cron=schedule.get("cron"),
                         time=schedule.get("time"),
+                        cron_offset=schedule.get("cron_offset"),
                         source=self,
                     ),
                 )

taskiq/schedule_sources/label_based.py

@@ -1,6 +1,6 @@
 from dataclasses import dataclass, field
-from datetime import datetime
-from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional
+from datetime import datetime, timedelta
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Union
 
 from taskiq.abc.broker import AsyncBroker
 from taskiq.kicker import AsyncKicker
@@ -21,6 +21,7 @@ class ScheduledTask:
     kwargs: Dict[str, Any]
     source: "ScheduleSource"  # Backward point to source which this task belongs to
     cron: Optional[str] = field(default=None)
+    cron_offset: Optional[Union[str, timedelta]] = field(default=None)
     time: Optional[datetime] = field(default=None)
 
     def __post_init__(self) -> None:

taskiq/scheduler/scheduler.py

@@ -0,0 +1,113 @@
+import datetime
+
+import pytz
+from tzlocal import get_localzone
+
+from taskiq.cli.scheduler.run import should_run
+from taskiq.schedule_sources.label_based import LabelScheduleSource
+from taskiq.scheduler.scheduler import ScheduledTask
+
+DUMMY_SOURCE = LabelScheduleSource(broker=None)  # type: ignore
+
+
+def test_should_run_success() -> None:
+    hour = datetime.datetime.utcnow().hour
+    assert should_run(
+        ScheduledTask(
+            task_name="",
+            labels={},
+            args=[],
+            kwargs={},
+            source=DUMMY_SOURCE,
+            cron=f"* {hour} * * *",
+        ),
+    )
+
+
+def test_should_run_cron_str_offset() -> None:
+    hour = datetime.datetime.now().hour
+    zone = get_localzone()
+    assert should_run(
+        ScheduledTask(
+            task_name="",
+            labels={},
+            args=[],
+            kwargs={},
+            source=DUMMY_SOURCE,
+            cron=f"* {hour} * * *",
+            cron_offset=str(zone),
+        ),
+    )
+
+
+def test_should_run_cron_td_offset() -> None:
+    offset = 2
+    hour = datetime.datetime.utcnow().hour + offset
+    assert should_run(
+        ScheduledTask(
+            task_name="",
+            labels={},
+            args=[],
+            kwargs={},
+            source=DUMMY_SOURCE,
+            cron=f"* {hour} * * *",
+            cron_offset=datetime.timedelta(hours=offset),
+        ),
+    )
+
+
+def test_time_utc_without_zone() -> None:
+    time = datetime.datetime.utcnow()
+    assert should_run(
+        ScheduledTask(
+            task_name="",
+            labels={},
+            args=[],
+            kwargs={},
+            source=DUMMY_SOURCE,
+            time=time - datetime.timedelta(seconds=1),
+        ),
+    )
+
+
+def test_time_utc_with_zone() -> None:
+    time = datetime.datetime.now(tz=pytz.UTC)
+    assert should_run(
+        ScheduledTask(
+            task_name="",
+            labels={},
+            args=[],
+            kwargs={},
+            source=DUMMY_SOURCE,
+            time=time - datetime.timedelta(seconds=1),
+        ),
+    )
+
+
+def test_time_utc_with_local_zone() -> None:
+    localtz = get_localzone()
+    time = datetime.datetime.now(tz=localtz)
+    assert should_run(
+        ScheduledTask(
+            task_name="",
+            labels={},
+            args=[],
+            kwargs={},
+            source=DUMMY_SOURCE,
+            time=time - datetime.timedelta(seconds=1),
+        ),
+    )
+
+
+def test_time_localtime_without_zone() -> None:
+    time = datetime.datetime.now(tz=pytz.FixedOffset(240)).replace(tzinfo=None)
+    assert not should_run(
+        ScheduledTask(
+            task_name="",
+            labels={},
+            args=[],
+            kwargs={},
+            source=DUMMY_SOURCE,
+            time=time - datetime.timedelta(seconds=1),
+        ),
+    )