Added datetime schedule. (#60)

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

Author: s3rius

taskiq/cli/scheduler/run.py

@@ -47,6 +47,20 @@ async def schedules_updater(
         await asyncio.sleep(scheduler.refresh_delay)
 
 
+def should_run(task: ScheduledTask) -> bool:
+    """
+    Checks if it's time to run a task.
+
+    :param task: task to check.
+    :return: True if task must be sent.
+    """
+    if task.cron is not None:
+        return is_now(task.cron)
+    if task.time is not None:
+        return task.time <= datetime.utcnow()
+    return False
+
+
 async def run_scheduler(args: SchedulerArgs) -> None:  # noqa: C901, WPS210, WPS213
     """
     Runs scheduler loop.
@@ -83,7 +97,7 @@ async def run_scheduler(args: SchedulerArgs) -> None:  # noqa: C901, WPS210, WPS
         not_fired_tasks = []
         for task in tasks:
             try:
-                ready = is_now(task.cron)
+                ready = should_run(task)
             except ValueError:
                 logger.warning(
                     "Cannot parse cron: %s for task: %s",

taskiq/cli/utils.py

@@ -4,7 +4,7 @@
 from importlib import import_module
 from logging import getLogger
 from pathlib import Path
-from typing import Any, Generator
+from typing import Any, Generator, List
 
 logger = getLogger("taskiq.worker")
 
@@ -51,7 +51,7 @@ def import_object(object_spec: str) -> Any:
     return getattr(module, import_spec[1])
 
 
-def import_from_modules(modules: list[str]) -> None:
+def import_from_modules(modules: List[str]) -> None:
     """
     Import all modules from modules variable.
 
@@ -66,7 +66,7 @@ def import_from_modules(modules: list[str]) -> None:
             logger.warning(f"Cannot import {module}")
 
 
-def import_tasks(modules: list[str], pattern: str, fs_discover: bool) -> None:
+def import_tasks(modules: List[str], pattern: str, fs_discover: bool) -> None:
     """
     Import tasks modules.
 

taskiq/scheduler/scheduler.py

@@ -1,5 +1,6 @@
-from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any, Callable, Dict, List
+from dataclasses import dataclass, field
+from datetime import datetime
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional
 
 from taskiq.abc.broker import AsyncBroker
 from taskiq.scheduler.merge_functions import preserve_all
@@ -16,7 +17,17 @@ class ScheduledTask:
     labels: Dict[str, Any]
     args: List[Any]
     kwargs: Dict[str, Any]
-    cron: str
+    cron: Optional[str] = field(default=None)
+    time: Optional[datetime] = field(default=None)
+
+    def __post_init__(self) -> None:
+        """
+        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.")
 
 
 class TaskiqScheduler:
@@ -44,4 +55,3 @@ async def startup(self) -> None:  # pragma: no cover
         Here you can do stuff, like creating
         connections or anything you'd like.
         """
-        await self.broker.startup()