Merge pull request #744 from jenstroeger/add-some-typing-broker

Further improve typing on external interface functions for Broker and Worker

Author: LincolnPuzey

dramatiq/broker.py

@@ -93,7 +93,7 @@ class Broker:
         overwrite when they are declared.
     """
 
-    def __init__(self, middleware: Optional[list[Middleware]] = None):
+    def __init__(self, middleware: Optional[list[Middleware]] = None) -> None:
         self.logger = get_logger(__name__, type(self))
         self.actors: dict[str, Actor] = {}
         self.queues: Any = {}  # Subclasses make this a set!
@@ -132,7 +132,7 @@ def add_middleware(
         *,
         before: Optional[type[Middleware]] = None,
         after: Optional[type[Middleware]] = None,
-    ):
+    ) -> None:
         """Add a middleware object to this broker.  The middleware is
         appended to the end of the middleware list by default.
 
@@ -229,7 +229,7 @@ def declare_queue(self, queue_name: str) -> None:  # pragma: no cover
         """
         raise NotImplementedError
 
-    def enqueue(self, message, *, delay: Optional[int] = None):  # pragma: no cover
+    def enqueue(self, message: Message, *, delay: Optional[int] = None) -> Message:  # pragma: no cover
         """Enqueue a message on this broker.
 
         Parameters:
@@ -241,7 +241,7 @@ def enqueue(self, message, *, delay: Optional[int] = None):  # pragma: no cover
         """
         raise NotImplementedError
 
-    def get_actor(self, actor_name: str):  # pragma: no cover
+    def get_actor(self, actor_name: str) -> Actor:  # pragma: no cover
         """Look up an actor by its name.
 
         Parameters:
@@ -337,7 +337,7 @@ def __iter__(self) -> Self:  # pragma: no cover
         """Returns this instance as a Message iterator."""
         return self
 
-    def ack(self, message: MessageProxy):  # pragma: no cover
+    def ack(self, message: MessageProxy) -> None:  # pragma: no cover
         """Acknowledge that a message has been processed, removing it
         from the broker.
 
@@ -346,15 +346,15 @@ def ack(self, message: MessageProxy):  # pragma: no cover
         """
         raise NotImplementedError
 
-    def nack(self, message: MessageProxy):  # pragma: no cover
+    def nack(self, message: MessageProxy) -> None:  # pragma: no cover
         """Move a message to the dead-letter queue.
 
         Parameters:
           message(MessageProxy): The message to reject.
         """
         raise NotImplementedError
 
-    def requeue(self, messages: Iterable[MessageProxy]):  # pragma: no cover
+    def requeue(self, messages: Iterable[MessageProxy]) -> None:  # pragma: no cover
         """Move unacked messages back to their queues.  This is called
         by consumer threads when they fail or are shut down.  The
         default implementation does nothing.
@@ -363,7 +363,7 @@ def requeue(self, messages: Iterable[MessageProxy]):  # pragma: no cover
           messages(Iterable[MessageProxy]): The messages to requeue.
         """
 
-    def __next__(self):  # pragma: no cover
+    def __next__(self) -> MessageProxy:  # pragma: no cover
         """Retrieve the next message off of the queue.  This method
         blocks until a message becomes available.
 
@@ -393,7 +393,7 @@ class MessageProxy:
         message_id: str
         message_timestamp: int
 
-    def __init__(self, message: Message):
+    def __init__(self, message: Message) -> None:
         self.failed = False
         self._message = message
         self._exception: Optional[BaseException] = None

dramatiq/worker.py

@@ -134,7 +134,7 @@ def resume(self) -> None:
         for child in self.workers:
             child.resume()
 
-    def stop(self, timeout: int = 600000):
+    def stop(self, timeout: int = 600000) -> None:
         """Gracefully stop the Worker and all of its consumers and
         workers.
 
@@ -235,15 +235,15 @@ def _add_worker(self) -> None:
 
 
 class _WorkerMiddleware(Middleware):
-    def __init__(self, worker: Worker):
+    def __init__(self, worker: Worker) -> None:
         self.logger = get_logger(__name__, type(self))
         self.worker = worker
 
-    def after_declare_queue(self, broker: Broker, queue_name: str):
+    def after_declare_queue(self, broker: Broker, queue_name: str) -> None:
         self.logger.debug("Adding consumer for queue %r.", queue_name)
         self.worker._add_consumer(queue_name)
 
-    def after_declare_delay_queue(self, broker: Broker, queue_name: str):
+    def after_declare_delay_queue(self, broker: Broker, queue_name: str) -> None:
         self.logger.debug("Adding consumer for delay queue %r.", queue_name)
         self.worker._add_consumer(queue_name, delay=True)
 
@@ -257,7 +257,7 @@ def __init__(
         prefetch: int,
         work_queue: PriorityQueue[tuple[int, MessageProxy]],
         worker_timeout: int,
-    ):
+    ) -> None:
         super().__init__(daemon=True)
 
         self.logger = get_logger(__name__, "ConsumerThread(%s)" % queue_name)
@@ -338,7 +338,7 @@ def handle_delayed_messages(self) -> None:
             self.post_process_message(message)
             self.delay_queue.task_done()
 
-    def handle_message(self, message: MessageProxy):
+    def handle_message(self, message: MessageProxy) -> None:
         """Handle a message received off of the underlying consumer.
         If the message has an eta, delay it.  Otherwise, put it on the
         work queue.
@@ -362,7 +362,7 @@ def handle_message(self, message: MessageProxy):
             message.fail()
             self.post_process_message(message)
 
-    def post_process_message(self, message: MessageProxy):
+    def post_process_message(self, message: MessageProxy) -> None:
         """Called by worker threads whenever they're done processing
         individual messages, signaling that each message is ready to
         be acked or rejected.
@@ -416,7 +416,7 @@ def post_process_message(self, message: MessageProxy):
 
                 return
 
-    def requeue_messages(self, messages: Iterable[MessageProxy]):
+    def requeue_messages(self, messages: Iterable[MessageProxy]) -> None:
         """Called on worker shutdown and whenever there is a
         connection error to move unacked messages back to their
         respective queues asap.
@@ -472,7 +472,7 @@ def __init__(
         consumers: dict[str, _ConsumerThread],
         work_queue: PriorityQueue[tuple[int, MessageProxy]],
         worker_timeout: int,
-    ):
+    ) -> None:
         super().__init__(daemon=True)
 
         self.logger = get_logger(__name__, "WorkerThread")
@@ -504,7 +504,7 @@ def run(self) -> None:
         self.broker.emit_before("worker_thread_shutdown", self)
         self.logger.debug("Worker thread stopped.")
 
-    def process_message(self, message: MessageProxy):
+    def process_message(self, message: MessageProxy) -> None:
         """Process a message pulled off of the work queue then push it
         back to its associated consumer for post processing. Stuff any SkipMessage
         exception or BaseException into the message [proxy] so that it may be used