wip

Author: cunla

pyproject.toml

@@ -63,6 +63,7 @@ dev = [
     "fakeredis~=2.28",
     "pyyaml>=6,<7",
     "mypy>=1.16.0",
+    "types-croniter>=6.0.0.20250411",
 ]
 
 [tool.hatch.build.targets.sdist]

scheduler/helpers/queues/queue_logic.py

@@ -30,6 +30,10 @@ class NoSuchJobError(Exception):
     pass
 
 
+class NoSuchRegistryError(Exception):
+    pass
+
+
 def perform_job(job_model: JobModel, connection: ConnectionType) -> Any:  # noqa
     """The main execution method. Invokes the job function with the job arguments.
 
@@ -155,11 +159,11 @@ def count(self) -> int:
             res += getattr(self, registry).count(connection=self.connection)
         return res
 
-    def get_registry(self, name: str) -> Union[None, JobNamesRegistry]:
+    def get_registry(self, name: str) -> JobNamesRegistry:
         name = name.lower()
         if name in Queue.REGISTRIES:
             return getattr(self, Queue.REGISTRIES[name])
-        return None
+        raise NoSuchRegistryError(f"Unknown registry name {name}")
 
     def get_all_job_names(self) -> List[str]:
         res = list()

scheduler/helpers/sentry_integration.py

@@ -52,8 +52,8 @@ def sentry_patched_perform_job(self: Any, job_model: JobModel, *args: Queue, **k
                     transaction.name = job_model.func_name
 
                 with sentry_sdk.start_transaction(
-                        transaction,
-                        custom_sampling_context={"rq_job": job_model},
+                    transaction,
+                    custom_sampling_context={"rq_job": job_model},
                 ):
                     rv = old_perform_job(self, job_model, *args, **kwargs)
 
@@ -70,11 +70,7 @@ def sentry_patched_perform_job(self: Any, job_model: JobModel, *args: Queue, **k
         old_handle_exception = Worker.handle_exception
 
         def sentry_patched_handle_exception(self: Worker, job: Any, *exc_info: Any, **kwargs: Any) -> Any:
-            retry = (
-                    hasattr(job, "retries_left")
-                    and job.retries_left
-                    and job.retries_left > 0
-            )
+            retry = hasattr(job, "retries_left") and job.retries_left and job.retries_left > 0
             failed = job._status == JobStatus.FAILED or job.is_failed
             if failed and not retry:
                 _capture_exception(exc_info)
@@ -89,9 +85,7 @@ def sentry_patched_handle_exception(self: Worker, job: Any, *exc_info: Any, **kw
         def sentry_patched_enqueue_job(self: Queue, job: Any, **kwargs: Any) -> Any:
             scope = sentry_sdk.get_current_scope()
             if scope.span is not None:
-                job.meta["_sentry_trace_headers"] = dict(
-                    scope.iter_trace_propagation_headers()
-                )
+                job.meta["_sentry_trace_headers"] = dict(scope.iter_trace_propagation_headers())
 
             return old_enqueue_job(self, job, **kwargs)
 

scheduler/management/commands/delete_failed_executions.py

@@ -1,4 +1,5 @@
 import click
+from django.core.management import CommandParser
 from django.core.management.base import BaseCommand
 
 from scheduler.helpers.queues import get_queue

scheduler/management/commands/import.py

@@ -52,13 +52,9 @@ def create_task_from_dict(task_dict: Dict[str, Any], update: bool) -> Optional[T
         if not settings.USE_TZ and not timezone.is_naive(target):
             target = timezone.make_naive(target)
         kwargs["scheduled_time"] = target
-    model_fields = set(map(
-        lambda field: field.attname,
-        filter(
-            lambda field: hasattr(field, "attname"),
-            Task._meta.get_fields()
-        )
-    ))
+    model_fields = set(
+        map(lambda field: field.attname, filter(lambda field: hasattr(field, "attname"), Task._meta.get_fields()))
+    )
     keys_to_ignore = list(filter(lambda _k: _k not in model_fields, kwargs.keys()))
     for k in keys_to_ignore:
         del kwargs[k]
@@ -119,7 +115,7 @@ def add_arguments(self, parser: CommandParser) -> None:
         )
 
     def handle(self, *args: Any, **options: Any) -> None:
-        file = open(options.get("filename")) if options.get("filename") else sys.stdin # type: ignore[arg-type]
+        file = open(options.get("filename")) if options.get("filename") else sys.stdin  # type: ignore[arg-type]
         jobs = list()
         if options.get("format") == "json":
             import json
@@ -143,4 +139,4 @@ def handle(self, *args: Any, **options: Any) -> None:
             Task.objects.all().delete()
 
         for job in jobs:
-            create_task_from_dict(job, update=options.get("update")) # type: ignore[arg-type]
+            create_task_from_dict(job, update=options.get("update"))  # type: ignore[arg-type]

scheduler/management/commands/run_job.py

@@ -27,7 +27,7 @@ def add_arguments(self, parser: CommandParser) -> None:
         )
         parser.add_argument("args", nargs="*", help="Args for callable")
 
-    def handle(self, **options:Any)-> None:
+    def handle(self, **options: Any) -> None:
         verbosity = int(options.get("verbosity", 1))
         timeout = options.get("timeout")
         result_ttl = options.get("result_ttl")

scheduler/models/args.py

@@ -1,5 +1,5 @@
 from datetime import datetime
-from typing import Callable
+from typing import Callable, Any, Tuple, Dict, Type
 
 from django.contrib.contenttypes.fields import GenericForeignKey
 from django.contrib.contenttypes.models import ContentType
@@ -9,7 +9,7 @@
 
 from scheduler.helpers import utils
 
-ARG_TYPE_TYPES_DICT = {
+ARG_TYPE_TYPES_DICT:Dict[str,Type] = {
     "str": str,
     "int": int,
     "bool": bool,
@@ -37,7 +37,7 @@ class ArgType(models.TextChoices):
     object_id = models.PositiveIntegerField()
     content_object = GenericForeignKey()
 
-    def clean(self):
+    def clean(self) -> None:
         if self.arg_type not in ARG_TYPE_TYPES_DICT:
             raise ValidationError(
                 {
@@ -61,15 +61,15 @@ def clean(self):
                 {"arg_type": ValidationError(_(f"Could not parse {self.val} as {self.arg_type}"), code="invalid")}
             )
 
-    def save(self, **kwargs):
+    def save(self, **kwargs: Any) -> None:
         super(BaseTaskArg, self).save(**kwargs)
         self.content_object.save()
 
-    def delete(self, **kwargs):
+    def delete(self, **kwargs: Any) -> None:
         super(BaseTaskArg, self).delete(**kwargs)
         self.content_object.save()
 
-    def value(self):
+    def value(self) -> Any:
         if self.arg_type == "callable":
             res = utils.callable_func(self.val)()
         elif self.arg_type == "datetime":
@@ -86,16 +86,16 @@ class Meta:
 
 
 class TaskArg(BaseTaskArg):
-    def __str__(self):
+    def __str__(self) -> str:
         return f"TaskArg[arg_type={self.arg_type},value={self.value()}]"
 
 
 class TaskKwarg(BaseTaskArg):
     key = models.CharField(max_length=255)
 
-    def __str__(self):
+    def __str__(self) -> str:
         key, value = self.value()
         return f"TaskKwarg[key={key},arg_type={self.arg_type},value={self.val}]"
 
-    def value(self):
+    def value(self) -> Tuple[str, Any]:
         return self.key, super(TaskKwarg, self).value()

scheduler/models/task.py

@@ -1,6 +1,6 @@
 import math
 from datetime import timedelta, datetime
-from typing import Dict, Any, Optional
+from typing import Dict, Any, Optional, List, Tuple, Callable
 
 import croniter
 from django.conf import settings as django_settings
@@ -29,11 +29,11 @@
 def _get_task_for_job(job: JobModel) -> Optional["Task"]:
     if job.task_type is None or job.scheduled_task_id is None:
         return None
-    task = Task.objects.filter(id=job.scheduled_task_id).first()
+    task: Task = Task.objects.filter(id=job.scheduled_task_id).first()
     return task
 
 
-def failure_callback(job: JobModel, connection, result, *args, **kwargs):
+def failure_callback(job: JobModel, connection:ConnectionType, result: Any, *args: Any, **kwargs: Any) -> None:
     task = _get_task_for_job(job)
     if task is None:
         logger.warn(f"Could not find task for job {job.name}")
@@ -48,7 +48,7 @@ def failure_callback(job: JobModel, connection, result, *args, **kwargs):
     task.save(schedule_job=True, clean=False)
 
 
-def success_callback(job: JobModel, connection: ConnectionType, result: Any, *args, **kwargs):
+def success_callback(job: JobModel, connection: ConnectionType, result: Any, *args: Any, **kwargs: Any) -> None:
     task = _get_task_for_job(job)
     if task is None:
         logger.warn(f"Could not find task for job {job.name}")
@@ -59,7 +59,7 @@ def success_callback(job: JobModel, connection: ConnectionType, result: Any, *ar
     task.save(schedule_job=True, clean=False)
 
 
-def get_queue_choices():
+def get_queue_choices() -> List[Tuple[str, str]]:
     queue_names = get_queue_names()
     return [(queue, queue) for queue in queue_names]
 
@@ -176,20 +176,20 @@ class TimeUnits(models.TextChoices):
         ),
     )
 
-    def callable_func(self):
+    def callable_func(self)->Callable:
         """Translate callable string to callable"""
         return utils.callable_func(self.callable)
 
-    @admin.display(boolean=True, description=_("is scheduled?"))
+    @admin.display(boolean=True, description=_("is scheduled?"))  # type: ignore[misc]
     def is_scheduled(self) -> bool:
         """Check whether a next job for this task is queued/scheduled to be executed"""
         if self.job_name is None:  # no job_id => is not scheduled
             return False
         # check whether job_id is in scheduled/queued/active jobs
         res = (
-            (self.job_name in self.rqueue.scheduled_job_registry.all())
-            or (self.job_name in self.rqueue.queued_job_registry.all())
-            or (self.job_name in self.rqueue.active_job_registry.all())
+                (self.job_name in self.rqueue.scheduled_job_registry.all())
+                or (self.job_name in self.rqueue.queued_job_registry.all())
+                or (self.job_name in self.rqueue.active_job_registry.all())
         )
         # If the job_id is not scheduled/queued/started,
         # update the job_id to None. (The job_id belongs to a previous run which is completed)
@@ -198,29 +198,29 @@ def is_scheduled(self) -> bool:
             super(Task, self).save()
         return res
 
-    @admin.display(description="Callable")
+    @admin.display(description="Callable")  # type: ignore[misc]
     def function_string(self) -> str:
         args = self.parse_args()
         args_list = [repr(arg) for arg in args]
         kwargs = self.parse_kwargs()
         kwargs_list = [k + "=" + repr(v) for (k, v) in kwargs.items()]
         return self.callable + f"({', '.join(args_list + kwargs_list)})"
 
-    def parse_args(self):
+    def parse_args(self) -> List[Any]:
         """Parse args for running the job"""
         args = self.callable_args.all()
         return [arg.value() for arg in args]
 
-    def parse_kwargs(self):
+    def parse_kwargs(self) -> Dict[str, Any]:
         """Parse kwargs for running the job"""
         kwargs = self.callable_kwargs.all()
         return dict([kwarg.value() for kwarg in kwargs])
 
-    def _next_job_id(self):
+    def _next_job_id(self) -> str:
         addition = timezone.now().strftime("%Y%m%d%H%M%S%f")
         return f"{self.queue}:{self.id}:{addition}"
 
-    def _enqueue_args(self) -> Dict:
+    def _enqueue_args(self) -> Dict[str, Any]:
         """Args for Queue.enqueue_call.
         Set all arguments for Queue.enqueue. Particularly:
         - set job timeout and ttl
@@ -282,7 +282,7 @@ def _schedule_time(self) -> datetime:
                 self.repeat = (self.repeat - gap) if self.repeat is not None else None
         return utc(self.scheduled_time) if django_settings.USE_TZ else self.scheduled_time
 
-    def to_dict(self) -> Dict:
+    def to_dict(self) -> Dict[str, Any]:
         """Export model to dictionary, so it can be saved as external file backup"""
         interval_unit = str(self.interval_unit) if self.interval_unit else None
         res = dict(
@@ -321,16 +321,11 @@ def to_dict(self) -> Dict:
         )
         return res
 
-    def get_absolute_url(self):
+    def get_absolute_url(self) -> str:
         model = self._meta.model.__name__.lower()
-        return reverse(
-            f"admin:scheduler_{model}_change",
-            args=[
-                self.id,
-            ],
-        )
+        return reverse(f"admin:scheduler_{model}_change", args=[self.id])
 
-    def __str__(self):
+    def __str__(self) -> str:
         func = self.function_string()
         return f"{self.task_type}[{self.name}={func}]"
 
@@ -359,7 +354,7 @@ def _schedule(self) -> bool:
         self.job_name = job.name
         return True
 
-    def save(self, **kwargs):
+    def save(self, **kwargs: Any) -> None:
         should_clean = kwargs.pop("clean", True)
         if should_clean:
             self.clean()
@@ -372,25 +367,25 @@ def save(self, **kwargs):
             self._schedule()
             super(Task, self).save()
 
-    def delete(self, **kwargs):
+    def delete(self, **kwargs: Any) -> None:
         self.unschedule()
         super(Task, self).delete(**kwargs)
 
-    def interval_seconds(self):
+    def interval_seconds(self) -> float:
         kwargs = {
             self.interval_unit: self.interval,
         }
         return timedelta(**kwargs).total_seconds()
 
-    def clean_callable(self):
+    def clean_callable(self) -> None:
         try:
             utils.callable_func(self.callable)
         except Exception:
             raise ValidationError(
                 {"callable": ValidationError(_("Invalid callable, must be importable"), code="invalid")}
             )
 
-    def clean_queue(self):
+    def clean_queue(self) -> None:
         queue_names = get_queue_names()
         if self.queue not in queue_names:
             raise ValidationError(
@@ -401,7 +396,7 @@ def clean_queue(self):
                 }
             )
 
-    def clean_interval_unit(self):
+    def clean_interval_unit(self) -> None:
         config = settings.SCHEDULER_CONFIG
         if config.SCHEDULER_INTERVAL > self.interval_seconds():
             raise ValidationError(
@@ -424,13 +419,13 @@ def clean_result_ttl(self) -> None:
                 params={"interval": self.interval_seconds()},
             )
 
-    def clean_cron_string(self):
+    def clean_cron_string(self) -> None:
         try:
             croniter.croniter(self.cron_string)
         except ValueError as e:
             raise ValidationError({"cron_string": ValidationError(_(str(e)), code="invalid")})
 
-    def clean(self):
+    def clean(self) -> None:
         if self.task_type not in TaskType.values:
             raise ValidationError(
                 {"task_type": ValidationError(_("Invalid task type"), code="invalid")},
@@ -470,7 +465,7 @@ def get_scheduled_task(task_type_str: str, task_id: int) -> Task:
             task = Task.objects.filter(task_type=task_type, id=task_id).first()
             if task is None:
                 raise ValueError(f"Job {task_type}:{task_id} does not exit")
-            return task
+            return task  # type: ignore[no-any-return]
         except ValueError:
             raise ValueError(f"Invalid task type {task_type_str}")
     raise ValueError(f"Job Model {task_type_str} does not exist, choices are {TASK_TYPES}")
@@ -484,5 +479,5 @@ def run_task(task_model: str, task_id: int) -> Any:
     logger.debug(f"Running task {str(scheduled_task)}")
     args = scheduled_task.parse_args()
     kwargs = scheduled_task.parse_kwargs()
-    res = scheduled_task.callable_func()(*args, **kwargs)
+    res = scheduled_task.callable_func()(*args, **kwargs)  # type: ignore[no-untyped-call]
     return res