make PriorityQueue and TriggeredService generic

Author: prandla

cms/io/priorityqueue.py

@@ -41,14 +41,13 @@
 from datetime import datetime
 from functools import total_ordering
 from typing import TypedDict
+import typing
 
 from gevent.event import Event
 
 from cmscommon.datetime import make_datetime, make_timestamp
 
 
-# TODO: make PriorityQueue generic over the exact type of the QueueItem.
-# Allows more exact typing in other classes (i.e. TriggeredService and its children).
 class QueueItem:
 
     """Payload of an item in the queue.
@@ -61,15 +60,16 @@ def to_dict(self) -> dict:
         """Return a dict() representation of the object."""
         return self.__dict__
 
+QueueItemT = typing.TypeVar('QueueItemT', bound=QueueItem)
 
 @total_ordering
-class QueueEntry:
+class QueueEntry(typing.Generic[QueueItemT]):
 
     """Type of the actual objects in the queue.
 
     """
 
-    def __init__(self, item: QueueItem, priority: int, timestamp: datetime, index: int):
+    def __init__(self, item: QueueItemT, priority: int, timestamp: datetime, index: int):
         """Create a QueueEntry object.
 
         item: the payload.
@@ -102,7 +102,7 @@ class QueueEntryDict(TypedDict):
     priority: int
     timestamp: float
 
-class PriorityQueue:
+class PriorityQueue(typing.Generic[QueueItemT]):
 
     """A priority queue.
 
@@ -125,11 +125,11 @@ def __init__(self):
         """Create a priority queue."""
         # The queue: a min-heap whose elements are of the form
         # (priority, timestamp, item), where item is the actual data.
-        self._queue: list[QueueEntry] = []
+        self._queue: list[QueueEntry[QueueItemT]] = []
 
         # Reverse lookup for the items in the queue: a dictionary
         # associating the index in the queue to each item.
-        self._reverse = {}
+        self._reverse: dict[QueueItemT, int] = {}
 
         # Event to signal that there are items in the queue.
         self._event = Event()
@@ -159,7 +159,7 @@ def _verify(self) -> bool:
                 return False
         return True
 
-    def __contains__(self, item: QueueItem) -> bool:
+    def __contains__(self, item: QueueItemT) -> bool:
         """Implement the 'in' operator for an item in the queue.
 
         item: an item to search.
@@ -234,7 +234,7 @@ def _updown_heap(self, idx: int) -> int:
         idx = self._up_heap(idx)
         return self._down_heap(idx)
 
-    def push(self, item: QueueItem, priority: int | None = None,
+    def push(self, item: QueueItemT, priority: int | None = None,
              timestamp: datetime | None = None) -> bool:
         """Push an item in the queue. If timestamp is not specified,
         uses the current time.
@@ -270,7 +270,7 @@ def push(self, item: QueueItem, priority: int | None = None,
 
         return True
 
-    def top(self, wait: bool = False) -> QueueEntry:
+    def top(self, wait: bool = False) -> QueueEntry[QueueItemT]:
         """Return the first element in the queue without extracting it.
 
         wait: if True, block until an element is present.
@@ -292,7 +292,7 @@ def top(self, wait: bool = False) -> QueueEntry:
                         continue
                     return self._queue[0]
 
-    def pop(self, wait: bool = False) -> QueueEntry:
+    def pop(self, wait: bool = False) -> QueueEntry[QueueItemT]:
         """Extract (and return) the first element in the queue.
 
         wait: if True, block until an element is present.
@@ -317,7 +317,7 @@ def pop(self, wait: bool = False) -> QueueEntry:
             self._event.clear()
         return top
 
-    def remove(self, item: QueueItem) -> QueueEntry:
+    def remove(self, item: QueueItemT) -> QueueEntry[QueueItemT]:
         """Remove an item from the queue. Raise a KeyError if not present.
 
         item: the item to remove.
@@ -343,7 +343,7 @@ def remove(self, item: QueueItem) -> QueueEntry:
 
         return entry
 
-    def set_priority(self, item: QueueItem, priority: int):
+    def set_priority(self, item: QueueItemT, priority: int):
         """Change the priority of an item inside the queue. Raises an
         exception if the item is not in the queue.
 

cms/io/triggeredservice.py

@@ -31,13 +31,13 @@
 from gevent.event import Event
 
 from cms.io import PriorityQueue, Service, rpc_method
-from cms.io.priorityqueue import QueueEntry, QueueEntryDict, QueueItem
+from cms.io.priorityqueue import QueueEntry, QueueEntryDict, QueueItemT
 
 
 logger = logging.getLogger(__name__)
 
 
-class Executor(metaclass=ABCMeta):
+class Executor(typing.Generic[QueueItemT], metaclass=ABCMeta):
 
     """A class taking care of executing operations.
 
@@ -61,9 +61,9 @@ def __init__(self, batch_executions: bool = False):
         super().__init__()
 
         self._batch_executions = batch_executions
-        self._operation_queue = PriorityQueue()
+        self._operation_queue: PriorityQueue[QueueItemT] = PriorityQueue()
 
-    def __contains__(self, item: QueueItem) -> bool:
+    def __contains__(self, item: QueueItemT) -> bool:
         """Return whether the item is in the queue.
 
         item: the item to look for.
@@ -84,7 +84,7 @@ def get_status(self) -> list[QueueEntryDict]:
         """
         return self._operation_queue.get_status()
 
-    def enqueue(self, item: QueueItem, priority: int | None = None, timestamp: datetime | None = None) -> bool:
+    def enqueue(self, item: QueueItemT, priority: int | None = None, timestamp: datetime | None = None) -> bool:
         """Add an item to the queue.
 
         item: the item to add.
@@ -97,7 +97,7 @@ def enqueue(self, item: QueueItem, priority: int | None = None, timestamp: datet
         """
         return self._operation_queue.push(item, priority, timestamp)
 
-    def dequeue(self, item: QueueItem) -> QueueEntry:
+    def dequeue(self, item: QueueItemT) -> QueueEntry[QueueItemT]:
         """Remove an item from the queue.
 
         item: the item to remove.
@@ -107,7 +107,7 @@ def dequeue(self, item: QueueItem) -> QueueEntry:
         """
         return self._operation_queue.remove(item)
 
-    def _pop(self, wait: bool = False) -> QueueEntry:
+    def _pop(self, wait: bool = False) -> QueueEntry[QueueItemT]:
         """Extract (and return) the first element in the queue.
 
         wait: if True, block until an element is present.
@@ -180,7 +180,7 @@ def max_operations_per_batch(self) -> int:
         return 0
 
     @abstractmethod
-    def execute(self, entry: QueueEntry | list[QueueEntry]):
+    def execute(self, entry: QueueEntry[QueueItemT] | list[QueueEntry[QueueItemT]]):
         """Perform a single operation.
 
         Must be implemented if batch_execution is false.
@@ -196,8 +196,10 @@ def execute(self, entry: QueueEntry | list[QueueEntry]):
         pass
 
 
+# The correct bound here is Executor[QueueItemT], but expressing that would
+# require higher-kinded types, which python's typechecking does not support.
 ExecutorT = typing.TypeVar('ExecutorT', bound=Executor)
-class TriggeredService(Service, typing.Generic[ExecutorT]):
+class TriggeredService(Service, typing.Generic[QueueItemT, ExecutorT]):
 
     """A service receiving notifications to perform an operation.
 
@@ -263,7 +265,7 @@ def get_executor(self) -> ExecutorT:
         """
         return self._executors[0]
 
-    def enqueue(self, operation: QueueItem, priority: int | None = None,
+    def enqueue(self, operation: QueueItemT, priority: int | None = None,
                 timestamp: datetime | None = None) -> int:
         """Add an operation to the queue of each executor.
 
@@ -282,7 +284,7 @@ def enqueue(self, operation: QueueItem, priority: int | None = None,
                 ret += 1
         return ret
 
-    def dequeue(self, operation: QueueItem):
+    def dequeue(self, operation: QueueItemT):
         """Remove an operation from the queue of each executor.
 
         operation: the operation to dequeue.

cms/service/EvaluationService.py

@@ -59,7 +59,7 @@
 logger = logging.getLogger(__name__)
 
 
-class EvaluationExecutor(Executor):
+class EvaluationExecutor(Executor[ESOperation]):
 
     # Real maximum number of operations to be sent to a worker.
     MAX_OPERATIONS_PER_BATCH = 25
@@ -122,7 +122,7 @@ def max_operations_per_batch(self) -> int:
                     ratio, ret)
         return ret
 
-    def execute(self, entries: list[QueueEntry]):
+    def execute(self, entries: list[QueueEntry[ESOperation]]):
         """Execute a batch of operations in the queue.
 
         The operations might not be executed immediately because of
@@ -190,7 +190,7 @@ def _pop(self, wait=False):
         self._remove_from_cumulative_status(queue_entry)
         return queue_entry
 
-    def _remove_from_cumulative_status(self, queue_entry: QueueEntry):
+    def _remove_from_cumulative_status(self, queue_entry: QueueEntry[ESOperation]):
         # Remove the item from the cumulative status dictionary.
         key = queue_entry.item.short_key() + (queue_entry.priority,)
         self.queue_status_cumulative[key]["item"]["multiplicity"] -= 1
@@ -223,7 +223,7 @@ def __init__(self, job: Job, job_success: bool):
         self.job_success = job_success
 
 
-class EvaluationService(TriggeredService[EvaluationExecutor]):
+class EvaluationService(TriggeredService[ESOperation, EvaluationExecutor]):
     """Evaluation service.
 
     """

cms/service/PrintingService.py

@@ -61,7 +61,7 @@ def to_dict(self):
         return {"printjob_id": self.printjob_id}
 
 
-class PrintingExecutor(Executor):
+class PrintingExecutor(Executor[PrintingOperation]):
     def __init__(self, file_cacher):
         super().__init__()
 
@@ -72,7 +72,7 @@ def __init__(self, file_cacher):
         self.jinja2_env.filters["escape_tex_normal"] = escape_tex_normal
         self.jinja2_env.filters["escape_tex_tt"] = escape_tex_tt
 
-    def execute(self, entry: QueueEntry):
+    def execute(self, entry: QueueEntry[PrintingOperation]):
         """Print a print job.
 
         This is the core of PrintingService.
@@ -208,12 +208,12 @@ def execute(self, entry: QueueEntry):
                 rmtree(directory)
 
 
-class PrintingService(TriggeredService):
+class PrintingService(TriggeredService[PrintingOperation, PrintingExecutor]):
     """A service that prepares print jobs and sends them to a printer.
 
     """
 
-    def __init__(self, shard):
+    def __init__(self, shard: int):
         """Initialize the PrintingService.
 
         """

cms/service/ProxyService.py

@@ -128,7 +128,7 @@ def to_dict(self):
                 "data": self.data}
 
 
-class ProxyExecutor(Executor):
+class ProxyExecutor(Executor[ProxyOperation]):
     """A thread that sends data to one ranking.
 
     The object is used as a thread-local storage and its run method is
@@ -182,7 +182,7 @@ def __init__(self, ranking: str):
         self._ranking = ranking
         self._visible_ranking = safe_url(ranking)
 
-    def execute(self, entries: list[QueueEntry]):
+    def execute(self, entries: list[QueueEntry[ProxyOperation]]):
         """Consume (i.e. send) the data put in the queue, forever.
 
         Pick all operations found in the queue (if there aren't any,
@@ -230,7 +230,7 @@ def execute(self, entries: list[QueueEntry]):
             gevent.sleep(self.FAILURE_WAIT)
 
 
-class ProxyService(TriggeredService):
+class ProxyService(TriggeredService[ProxyOperation, ProxyExecutor]):
     """Maintain the information held by rankings up-to-date.
 
     Discover (by receiving notifications and by periodically sweeping

cms/service/ScoringService.py

@@ -38,12 +38,12 @@
 logger = logging.getLogger(__name__)
 
 
-class ScoringExecutor(Executor):
+class ScoringExecutor(Executor[ScoringOperation]):
     def __init__(self, proxy_service):
         super().__init__()
         self.proxy_service = proxy_service
 
-    def execute(self, entry: QueueEntry):
+    def execute(self, entry: QueueEntry[ScoringOperation]):
         """Assign a score to a submission result.
 
         This is the core of ScoringService: here we retrieve the result
@@ -113,7 +113,7 @@ def execute(self, entry: QueueEntry):
                     submission_id=submission.id)
 
 
-class ScoringService(TriggeredService):
+class ScoringService(TriggeredService[ScoringOperation, ScoringExecutor]):
     """A service that assigns a score to submission results.
 
     A submission result is ready to be scored when its compilation is

cmstestsuite/unit_tests/io/triggeredservice_test.py

@@ -40,7 +40,7 @@ def get_notifications(self):
         return self._notifications
 
 
-class FakeExecutor(Executor):
+class FakeExecutor(Executor[FakeQueueItem]):
     def __init__(self, notifier, batch_executions=False):
         super().__init__(batch_executions)
         self._notifier = notifier