Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions taskiq/scheduler/scheduled_task/cron_spec.py
Original file line number Diff line number Diff line change
@@ -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."""
Expand All @@ -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."""
Expand Down
7 changes: 5 additions & 2 deletions taskiq/scheduler/scheduled_task/task.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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

Expand Down
24 changes: 24 additions & 0 deletions taskiq/scheduler/scheduled_task/validators.py
Original file line number Diff line number Diff line change
@@ -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(
Expand Down
39 changes: 38 additions & 1 deletion tests/scheduler/test_scheduled_task.py
Original file line number Diff line number Diff line change
@@ -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:
Expand Down
Loading