Improve typing of middleware functions

Author: jenstroeger

dramatiq/middleware/age_limit.py

@@ -17,6 +17,8 @@
 
 from __future__ import annotations
 
+from typing import Optional
+
 from ..common import current_millis
 from ..logging import get_logger
 from .middleware import Middleware, SkipMessage
@@ -32,7 +34,7 @@ class AgeLimit(Middleware):
         indefinitely.
     """
 
-    def __init__(self, *, max_age=None):
+    def __init__(self, *, max_age: Optional[int] = None) -> None:
         self.logger = get_logger(__name__, type(self))
         self.max_age = max_age
 

dramatiq/middleware/asyncio.py

@@ -25,7 +25,7 @@
 class AsyncIO(Middleware):
     """This middleware manages the event loop thread for async actors."""
 
-    def __init__(self):
+    def __init__(self) -> None:
         self.logger = get_logger(__name__, type(self))
 
     def before_worker_boot(self, broker, worker):

dramatiq/middleware/group_callbacks.py

@@ -19,14 +19,14 @@
 
 import os
 
-from ..rate_limits import Barrier
+from ..rate_limits import Barrier, RateLimiterBackend
 from .middleware import Middleware
 
 GROUP_CALLBACK_BARRIER_TTL = int(os.getenv("dramatiq_group_callback_barrier_ttl", "86400000"))
 
 
 class GroupCallbacks(Middleware):
-    def __init__(self, rate_limiter_backend):
+    def __init__(self, rate_limiter_backend: RateLimiterBackend) -> None:
         self.rate_limiter_backend = rate_limiter_backend
 
     def after_process_message(self, broker, message, *, result=None, exception=None):

dramatiq/middleware/middleware.py

@@ -18,6 +18,14 @@
 
 from __future__ import annotations
 
+from typing import TYPE_CHECKING, Any, Callable, Optional
+
+if TYPE_CHECKING:
+    from ..actor import Actor
+    from ..broker import Broker, MessageProxy
+    from ..message import Message
+    from ..worker import Worker, _ConsumerThread, _WorkerThread
+
 
 class MiddlewareError(Exception):
     """Base class for middleware errors."""
@@ -36,39 +44,39 @@ class Middleware:
     """
 
     @property
-    def actor_options(self):
+    def actor_options(self) -> set[str]:
         """The set of options that may be configured on each actor."""
         return set()
 
     @property
-    def forks(self):
+    def forks(self) -> list[Callable[[], int]]:
         """A list of functions to run in separate forks of the main
         process.
         """
         return []
 
-    def before_ack(self, broker, message):
+    def before_ack(self, broker: Broker, message: MessageProxy) -> None:
         """Called before a message is acknowledged."""
 
-    def after_ack(self, broker, message):
+    def after_ack(self, broker: Broker, message: MessageProxy) -> None:
         """Called after a message has been acknowledged."""
 
-    def before_nack(self, broker, message):
+    def before_nack(self, broker: Broker, message: MessageProxy) -> None:
         """Called before a message is rejected."""
 
-    def after_nack(self, broker, message):
+    def after_nack(self, broker: Broker, message: MessageProxy) -> None:
         """Called after a message has been rejected."""
 
-    def before_declare_actor(self, broker, actor):
+    def before_declare_actor(self, broker: Broker, actor: Actor) -> None:
         """Called before an actor is declared."""
 
-    def after_declare_actor(self, broker, actor):
+    def after_declare_actor(self, broker: Broker, actor: Actor) -> None:
         """Called after an actor has been declared."""
 
-    def before_declare_queue(self, broker, queue_name):
+    def before_declare_queue(self, broker: Broker, queue_name: str) -> None:
         """Called before a queue is declared."""
 
-    def after_declare_queue(self, broker, queue_name):
+    def after_declare_queue(self, broker: Broker, queue_name: str) -> None:
         """Called after a queue has been declared.
 
         This signals that the queue has been registered with the
@@ -78,19 +86,19 @@ def after_declare_queue(self, broker, queue_name):
         them until messages are enqueued or consumed.
         """
 
-    def after_declare_delay_queue(self, broker, queue_name):
+    def after_declare_delay_queue(self, broker: Broker, queue_name: str) -> None:
         """Called after a delay queue has been declared."""
 
-    def before_enqueue(self, broker, message, delay):
+    def before_enqueue(self, broker: Broker, message: Message, delay: int) -> None:
         """Called before a message is enqueued."""
 
-    def after_enqueue(self, broker, message, delay):
+    def after_enqueue(self, broker: Broker, message: Message, delay: int) -> None:
         """Called after a message has been enqueued."""
 
-    def before_delay_message(self, broker, message):
+    def before_delay_message(self, broker: Broker, message: MessageProxy) -> None:
         """Called before a message has been delayed in worker memory."""
 
-    def before_process_message(self, broker, message):
+    def before_process_message(self, broker: Broker, message: MessageProxy) -> None:
         """Called before a message is processed.
 
         Raises:
@@ -99,42 +107,49 @@ def before_process_message(self, broker, message):
             of ``after_process_message``.
         """
 
-    def after_process_message(self, broker, message, *, result=None, exception=None):
+    def after_process_message(
+        self,
+        broker: Broker,
+        message: MessageProxy,
+        *,
+        result: Optional[Any] = None,
+        exception: Optional[BaseException] = None,
+    ) -> None:
         """Called after a message has been processed."""
 
-    def after_skip_message(self, broker, message):
+    def after_skip_message(self, broker: Broker, message: MessageProxy) -> None:
         """Called instead of ``after_process_message`` after a message
         has been skipped.
         """
 
-    def after_process_boot(self, broker):
+    def after_process_boot(self, broker: Broker) -> None:
         """Called immediately after subprocess start up."""
 
-    def before_worker_boot(self, broker, worker):
+    def before_worker_boot(self, broker: Broker, worker: Worker) -> None:
         """Called before the worker process starts up."""
 
-    def after_worker_boot(self, broker, worker):
+    def after_worker_boot(self, broker: Broker, worker: Worker) -> None:
         """Called after the worker process has started up."""
 
-    def before_worker_shutdown(self, broker, worker):
+    def before_worker_shutdown(self, broker: Broker, worker: Worker) -> None:
         """Called before the worker process shuts down."""
 
-    def after_worker_shutdown(self, broker, worker):
+    def after_worker_shutdown(self, broker: Broker, worker: Worker) -> None:
         """Called after the worker process shuts down."""
 
-    def after_consumer_thread_boot(self, broker, thread):
+    def after_consumer_thread_boot(self, broker: Broker, thread: _ConsumerThread) -> None:
         """Called from a consumer thread after it starts but before it starts its run loop."""
 
-    def before_consumer_thread_shutdown(self, broker, thread):
+    def before_consumer_thread_shutdown(self, broker: Broker, thread: _ConsumerThread) -> None:
         """Called before a consumer thread shuts down.  This may be
         used to clean up thread-local resources (such as Django
         database connections).
         """
 
-    def after_worker_thread_boot(self, broker, thread):
+    def after_worker_thread_boot(self, broker: Broker, thread: _WorkerThread) -> None:
         """Called from a worker thread after it starts but before it starts its run loop."""
 
-    def before_worker_thread_shutdown(self, broker, thread):
+    def before_worker_thread_shutdown(self, broker: Broker, thread: _WorkerThread) -> None:
         """Called before a worker thread shuts down.  This may be used
         to clean up thread-local resources (such as Django database
         connections).

dramatiq/middleware/prometheus.py

@@ -45,10 +45,10 @@ class Prometheus(Middleware):
     .. _Prometheus: https://prometheus.io
     """
 
-    def __init__(self):
+    def __init__(self) -> None:
         self.logger = get_logger(__name__, type(self))
-        self.delayed_messages = set()
-        self.message_start_times = {}
+        self.delayed_messages: set[str] = set()
+        self.message_start_times: dict[str, int] = {}
 
     @property
     def forks(self):

dramatiq/middleware/retries.py

@@ -19,6 +19,7 @@
 
 import time
 import traceback
+from typing import Callable, Optional
 
 from ..common import compute_backoff
 from ..errors import Retry
@@ -60,15 +61,22 @@ class Retries(Middleware):
         apply to retried tasks.  Defaults to 15 seconds.
       max_backoff(int): The maximum amount of backoff milliseconds to
         apply to retried tasks.  Defaults to 7 days.
-      retry_when(Callable[[int, Exception], bool]): An optional
+      retry_when(Callable[[int, BaseException], bool]): An optional
         predicate that can be used to programmatically determine
         whether a task should be retried or not.  This takes
         precedence over `max_retries` when set.
       on_retry_exhausted(str): Name of an actor to send a message to when
        message is failed due to retries being exceeded.
     """
 
-    def __init__(self, *, max_retries=20, min_backoff=None, max_backoff=None, retry_when=None):
+    def __init__(
+        self,
+        *,
+        max_retries: int = 20,
+        min_backoff: Optional[int] = None,
+        max_backoff: Optional[int] = None,
+        retry_when: Optional[Callable[[int, BaseException], bool]] = None,
+    ) -> None:
         self.logger = get_logger(__name__, type(self))
         self.max_retries = max_retries
         self.min_backoff = min_backoff or DEFAULT_MIN_BACKOFF

dramatiq/middleware/shutdown.py

@@ -19,7 +19,6 @@
 
 import threading
 import warnings
-from typing import Optional, Type
 
 from ..logging import get_logger
 from .middleware import Middleware
@@ -55,9 +54,11 @@ class ShutdownNotifications(Middleware):
         Defaults to False, meaning actors will not be interrupted, and allowed to finish.
     """
 
-    def __init__(self, notify_shutdown=False):
+    def __init__(self, notify_shutdown: bool = False) -> None:
         self.logger = get_logger(__name__, type(self))
         self.notify_shutdown = notify_shutdown
+
+        self.manager: _ShutdownManager
         if is_gevent_active():
             self.manager = _GeventShutdownManager(self.logger)
         else:
@@ -132,11 +133,10 @@ def shutdown(self):
             raise_thread_exception(thread_id, Shutdown)
 
 
-_GeventShutdownManager: Optional[Type[_ShutdownManager]] = None
 if is_gevent_active():
     from gevent import getcurrent
 
-    class __GeventShutdownManager(_ShutdownManager):
+    class _GeventShutdownManager(_ShutdownManager):
 
         def __init__(self, logger=None):
             self.logger = logger or get_logger(__name__, type(self))
@@ -161,5 +161,3 @@ def shutdown(self):
                     thread_id,
                 )
                 greenlet.kill(Shutdown, block=False)
-
-    _GeventShutdownManager = __GeventShutdownManager

dramatiq/middleware/time_limit.py

@@ -21,7 +21,7 @@
 import warnings
 from threading import Thread
 from time import monotonic, sleep
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING, Optional, Union
 
 from ..logging import get_logger
 from .middleware import Middleware
@@ -64,10 +64,11 @@ class TimeLimit(Middleware):
         Defaults to 1 second (1,000 milliseconds).
     """
 
-    def __init__(self, *, time_limit=600000, interval=1000):
+    def __init__(self, *, time_limit: float = 600000, interval: int = 1000) -> None:
         self.logger = get_logger(__name__, type(self))
         self.time_limit = time_limit
 
+        self.manager: Union[_GeventTimeoutManager, _CtypesTimeoutManager]
         if is_gevent_active():
             self.manager = _GeventTimeoutManager(logger=self.logger)
         else:
@@ -162,7 +163,7 @@ def remove_timeout(self, thread_id):
             timer.close()
 
 
-_GeventTimeout: Optional["gevent.Timeout"] = None
+_GeventTimeout: Optional[gevent.Timeout] = None
 if is_gevent_active():
     from gevent import Timeout