fix: round-trip timedelta cron offset

Author: danfimov

taskiq/scheduler/scheduled_task/cron_spec.py

@@ -1,7 +1,7 @@
-from datetime import timedelta
-
 from pydantic import BaseModel
 
+from taskiq.scheduler.scheduled_task.validators import CronOffset
+
 
 class CronSpec(BaseModel):
     """Cron specification for running tasks."""
@@ -12,7 +12,7 @@ class CronSpec(BaseModel):
     months: str | int | None = "*"
     weekdays: str | int | None = "*"
 
-    offset: str | timedelta | None = None
+    offset: CronOffset | None = None
 
     def to_cron(self) -> str:  # pragma: no cover
         """Converts cron spec to cron string."""

taskiq/scheduler/scheduled_task/task.py

@@ -5,7 +5,10 @@
 
 from pydantic import BaseModel, Field, model_validator
 
-from taskiq.scheduler.scheduled_task.validators import validate_interval_value
+from taskiq.scheduler.scheduled_task.validators import (
+    CronOffset,
+    validate_interval_value,
+)
 
 if sys.version_info >= (3, 11):
     from typing import Self
@@ -23,7 +26,7 @@ class ScheduledTask(BaseModel):
     task_id: str | None = None
     schedule_id: str = Field(default_factory=lambda: uuid.uuid4().hex)
     cron: str | None = None
-    cron_offset: str | timedelta | None = None
+    cron_offset: CronOffset | None = None
     time: datetime | None = None
     interval: int | timedelta | None = None
 

taskiq/scheduler/scheduled_task/validators.py

@@ -1,4 +1,28 @@
 from datetime import timedelta
+from typing import Annotated, Any
+
+from pydantic import Discriminator, Tag, ValidationError
+
+from taskiq.compat import parse_obj_as
+
+
+def discriminate_offset(value: Any) -> str:
+    """Pick the union branch a raw cron offset value belongs to."""
+    if isinstance(value, timedelta):
+        return "timedelta"
+    if isinstance(value, str):
+        try:
+            parse_obj_as(timedelta, value)
+        except ValidationError:
+            return "str"
+        return "timedelta"
+    return "str"
+
+
+CronOffset = Annotated[
+    Annotated[timedelta, Tag("timedelta")] | Annotated[str, Tag("str")],
+    Discriminator(discriminate_offset),
+]
 
 
 def validate_interval_value(

tests/scheduler/test_scheduled_task.py

@@ -1,6 +1,43 @@
+from datetime import timedelta
+from typing import Any
+
 import pytest
 
-from taskiq.scheduler.scheduled_task import ScheduledTask
+from taskiq.scheduler.scheduled_task import CronSpec, ScheduledTask
+
+
+@pytest.mark.parametrize(
+    ("offset", "offset_type"),
+    [
+        (timedelta(hours=4), timedelta),
+        ("US/Eastern", str),
+    ],
+)
+def test_cron_spec_offset_roundtrip(offset: Any, offset_type: type) -> None:
+    restored = CronSpec.model_validate(CronSpec(offset=offset).model_dump(mode="json"))
+    assert restored.offset == offset
+    assert type(restored.offset) is offset_type
+
+
+@pytest.mark.parametrize(
+    ("offset", "offset_type"),
+    [
+        (timedelta(hours=2), timedelta),
+        ("US/Eastern", str),
+    ],
+)
+def test_scheduled_task_cron_offset_roundtrip(offset: Any, offset_type: type) -> None:
+    task = ScheduledTask(
+        task_name="a",
+        labels={},
+        args=[],
+        kwargs={},
+        cron="* * * * *",
+        cron_offset=offset,
+    )
+    restored = ScheduledTask.model_validate(task.model_dump(mode="json"))
+    assert restored.cron_offset == offset
+    assert type(restored.cron_offset) is offset_type
 
 
 def test_scheduled_task_parameters() -> None: