Improve typing

Author: emyller

src/task_processor/managers.py

@@ -3,21 +3,16 @@
 from django.db.models import Manager
 
 if typing.TYPE_CHECKING:
+    from django.db.models.query import RawQuerySet
+
     from task_processor.models import RecurringTask, Task
 
 
 class TaskManager(Manager["Task"]):
-    def get_tasks_to_process(self, num_tasks: int) -> typing.List["Task"]:
-        return list(
-            self.raw(
-                "SELECT * FROM get_tasks_to_process(%s)",
-                [num_tasks],
-            ),
-        )
+    def get_tasks_to_process(self, num_tasks: int) -> "RawQuerySet[Task]":
+        return self.raw("SELECT * FROM get_tasks_to_process(%s)", [num_tasks])
 
 
 class RecurringTaskManager(Manager["RecurringTask"]):
-    def get_tasks_to_process(self) -> typing.List["RecurringTask"]:
-        return list(
-            self.raw("SELECT * FROM get_recurringtasks_to_process()"),
-        )
+    def get_tasks_to_process(self) -> "RawQuerySet[RecurringTask]":
+        return self.raw("SELECT * FROM get_recurringtasks_to_process()")

src/task_processor/processor.py

@@ -9,7 +9,7 @@
 from django.utils import timezone
 
 from task_processor import metrics
-from task_processor.managers import TaskManager
+from task_processor.managers import RecurringTaskManager, TaskManager
 from task_processor.models import (
     AbstractBaseTask,
     RecurringTask,
@@ -68,7 +68,8 @@ def run_recurring_tasks(database: str) -> list[RecurringTaskRun]:
     # NOTE: We will probably see a lot of delay in the execution of recurring tasks
     # if the tasks take longer then `run_every` to execute. This is not
     # a problem for now, but we should be mindful of this limitation
-    tasks = RecurringTask.objects.db_manager(database).get_tasks_to_process()
+    task_manager: RecurringTaskManager = RecurringTask.objects.db_manager(database)
+    tasks = task_manager.get_tasks_to_process()
     if tasks:
         logger.debug(f"Running {len(tasks)} recurring task(s)")