chore: expose ConsumerThread and WorkerThread from worker module

Author: synweap15

dramatiq/__init__.py

@@ -38,7 +38,7 @@
 from .logging import get_logger
 from .message import Message, get_encoder, set_encoder
 from .middleware import Middleware
-from .worker import Worker
+from .worker import ConsumerThread, Worker, WorkerThread
 
 __all__ = [
     # Actors
@@ -80,6 +80,8 @@
     "Middleware",
     # Workers
     "Worker",
+    "ConsumerThread",
+    "WorkerThread",
 ]
 
 __version__ = "1.18.0"

dramatiq/middleware/middleware.py

@@ -24,7 +24,7 @@
     from ..actor import Actor
     from ..broker import Broker, MessageProxy
     from ..message import Message
-    from ..worker import Worker, _ConsumerThread, _WorkerThread
+    from ..worker import ConsumerThread, Worker, WorkerThread
 
 
 class MiddlewareError(Exception):
@@ -137,19 +137,19 @@ def before_worker_shutdown(self, broker: Broker, worker: Worker) -> None:
     def after_worker_shutdown(self, broker: Broker, worker: Worker) -> None:
         """Called after the worker process shuts down."""
 
-    def after_consumer_thread_boot(self, broker: Broker, thread: _ConsumerThread) -> None:
+    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: Broker, thread: _ConsumerThread) -> None:
+    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: Broker, thread: _WorkerThread) -> None:
+    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: Broker, thread: _WorkerThread) -> None:
+    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/worker.py

@@ -600,3 +600,7 @@ def stop(self) -> None:
 @lru_cache(maxsize=128)
 def has_results_middleware(broker: Broker) -> bool:
     return any(type(m) is Results for m in broker.middleware)
+
+
+ConsumerThread = _ConsumerThread
+WorkerThread = _WorkerThread