Fixed scheduled_task.

Signed-off-by: Pavel Kirilin <[email protected]>

.github/workflows/test.yml

@@ -37,6 +37,7 @@ jobs:
     strategy:
       matrix:
         py_version: ["3.8", "3.9", "3.10", "3.11", "3.12"]
+        pydantic_ver: ["<2", ">=2,<3"]
         os: [ubuntu-latest, windows-latest]
     runs-on: "${{ matrix.os }}"
     steps:
@@ -50,6 +51,8 @@ jobs:
           cache: "poetry"
       - name: Install deps
         run: poetry install --all-extras
+      - name: Setup pydantic version
+        run: poetry run pip install "pydantic ${{ matrix.pydantic_ver }}"
       - name: Run pytest check
         run: poetry run pytest -vv -n auto --cov="taskiq" .
       - name: Generate report

taskiq/result/__init__.py

@@ -1,12 +1,9 @@
-# flake8: noqa
-from packaging.version import Version
+from taskiq.compat import IS_PYDANTIC2
 
-from taskiq.compat import PYDANTIC_VER
-
-if PYDANTIC_VER >= Version("2.0"):
+if IS_PYDANTIC2:
     from .v2 import TaskiqResult
 else:
-    from .v1 import TaskiqResult
+    from .v1 import TaskiqResult  # type: ignore
 
 
 __all__ = [

taskiq/scheduler/scheduled_task/__init__.py

@@ -0,0 +1,11 @@
+from taskiq.compat import IS_PYDANTIC2
+
+from .cron_spec import CronSpec
+
+if IS_PYDANTIC2:
+    from .v2 import ScheduledTask
+else:
+    from .v1 import ScheduledTask  # type: ignore
+
+
+__all__ = ["CronSpec", "ScheduledTask"]

taskiq/scheduler/scheduled_task/cron_spec.py

@@ -0,0 +1,20 @@
+from datetime import timedelta
+from typing import Optional, Union
+
+from pydantic import BaseModel
+
+
+class CronSpec(BaseModel):
+    """Cron specification for running tasks."""
+
+    minutes: Optional[Union[str, int]] = "*"
+    hours: Optional[Union[str, int]] = "*"
+    days: Optional[Union[str, int]] = "*"
+    months: Optional[Union[str, int]] = "*"
+    weekdays: Optional[Union[str, int]] = "*"
+
+    offset: Optional[Union[str, timedelta]] = None
+
+    def to_cron(self) -> str:  # pragma: no cover
+        """Converts cron spec to cron string."""
+        return f"{self.minutes} {self.hours} {self.days} {self.months} {self.weekdays}"

taskiq/scheduler/scheduled_task/v1.py

@@ -0,0 +1,30 @@
+import uuid
+from datetime import datetime, timedelta
+from typing import Any, Dict, List, Optional, Union
+
+from pydantic import BaseModel, Field, root_validator
+
+
+class ScheduledTask(BaseModel):
+    """Abstraction over task schedule."""
+
+    task_name: str
+    labels: Dict[str, Any]
+    args: List[Any]
+    kwargs: Dict[str, Any]
+    schedule_id: str = Field(default_factory=lambda: uuid.uuid4().hex)
+    cron: Optional[str] = None
+    cron_offset: Optional[Union[str, timedelta]] = None
+    time: Optional[datetime] = None
+
+    @root_validator(pre=False)  # type: ignore
+    @classmethod
+    def __check(cls, values: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        This method validates, that either `cron` or `time` field is present.
+
+        :raises ValueError: if cron and time are none.
+        """
+        if values.get("cron") is None and values.get("time") is None:
+            raise ValueError("Either cron or datetime must be present.")
+        return values

taskiq/scheduler/scheduled_task/v2.py

@@ -0,0 +1,30 @@
+import uuid
+from datetime import datetime, timedelta
+from typing import Any, Dict, List, Optional, Union
+
+from pydantic import BaseModel, Field, model_validator
+from typing_extensions import Self
+
+
+class ScheduledTask(BaseModel):
+    """Abstraction over task schedule."""
+
+    task_name: str
+    labels: Dict[str, Any]
+    args: List[Any]
+    kwargs: Dict[str, Any]
+    schedule_id: str = Field(default_factory=lambda: uuid.uuid4().hex)
+    cron: Optional[str] = None
+    cron_offset: Optional[Union[str, timedelta]] = None
+    time: Optional[datetime] = None
+
+    @model_validator(mode="after")
+    def __check(self) -> Self:
+        """
+        This method validates, that either `cron` or `time` field is present.
+
+        :raises ValueError: if cron and time are none.
+        """
+        if self.cron is None and self.time is None:
+            raise ValueError("Either cron or datetime must be present.")
+        return self

tests/formatters/test_json_formatter.py

@@ -1,3 +1,5 @@
+import json
+
 import pytest
 
 from taskiq.formatters.json_formatter import JSONFormatter
@@ -24,7 +26,11 @@ async def test_json_dumps() -> None:
         ),
         labels={"label1": 1, "label2": "text"},
     )
-    assert fmt.dumps(msg) == expected
+    dumped = fmt.dumps(msg)
+    assert dumped.task_id == expected.task_id
+    assert dumped.task_name == expected.task_name
+    assert dumped.labels == expected.labels
+    assert json.loads(dumped.message) == json.loads(expected.message)
 
 
 @pytest.mark.anyio

tests/scheduler/test_scheduled_task.py

@@ -0,0 +1,14 @@
+import pytest
+
+from taskiq.scheduler.scheduled_task import ScheduledTask
+
+
+def test_scheduled_task_paramters() -> None:
+    with pytest.raises(ValueError):
+        ScheduledTask(
+            task_name="a",
+            labels={},
+            args=[],
+            kwargs={},
+            schedule_id="b",
+        )