Merge branch 'main' into tracer-middleware

Author: ndevenish

src/workflows/contrib/start_service.py

@@ -1,12 +1,14 @@
 from __future__ import annotations
 
 import sys
+from collections.abc import Callable
 from optparse import SUPPRESS_HELP, OptionParser
 
 import workflows
 import workflows.frontend
 import workflows.services
 import workflows.transport
+from workflows.transport.common_transport import CommonTransport
 
 
 class ServiceStarter:
@@ -27,12 +29,14 @@ def on_parsing(options, args):
         """
 
     @staticmethod
-    def on_transport_factory_preparation(transport_factory):
+    def on_transport_factory_preparation(
+        transport_factory,
+    ) -> Callable[[], CommonTransport] | None:
         """Plugin hook to intercept/manipulate newly created Transport factories
         before first invocation."""
 
     @staticmethod
-    def on_transport_preparation(transport):
+    def on_transport_preparation(transport: CommonTransport) -> CommonTransport | None:
         """Plugin hook to intercept/manipulate newly created Transport objects
         before connecting."""
 
@@ -136,7 +140,9 @@ def run(
             parser.error(f"Please specify a service name. {known_services_help}")
 
         # Create Transport factory
-        transport_factory = workflows.transport.lookup(options.transport)
+        transport_factory: Callable[[], CommonTransport] = workflows.transport.lookup(
+            options.transport
+        )
 
         # Call on_transport_factory_preparation hook
         transport_factory = (
@@ -147,7 +153,7 @@ def run(
         # Set up on_transport_preparation hook to affect newly created transport objects
         true_transport_factory = transport_factory
 
-        def on_transport_preparation_hook():
+        def on_transport_preparation_hook() -> CommonTransport:
             transport_object = true_transport_factory()
             return self.on_transport_preparation(transport_object) or transport_object
 

src/workflows/frontend/__init__.py

@@ -4,13 +4,15 @@
 import multiprocessing
 import threading
 import time
+from collections.abc import Callable
 
 import workflows
 import workflows.frontend.utilization
 import workflows.services
 import workflows.transport
 import workflows.util
 from workflows.services.common_service import CommonService
+from workflows.transport.common_transport import CommonTransport
 
 basestring = (str, bytes)
 
@@ -22,9 +24,11 @@ class Frontend:
     service.
     """
 
+    _transport_factory: Callable[[], CommonTransport]
+
     def __init__(
         self,
-        transport=None,
+        transport: Callable[[], CommonTransport] | str | None = None,
         service=None,
         transport_command_channel=None,
         restart_service=False,
@@ -98,7 +102,6 @@ def __getitem__(self, key):
         self.log = logging.LoggerAdapter(
             logging.getLogger("workflows.frontend"), LogAdapter()
         )
-        self.log.warn = self.log.warning  # Add support for deprecated .warn
 
         # Connect to the network transport layer
         if transport is None or isinstance(transport, basestring):

src/workflows/transport/__init__.py

@@ -11,11 +11,12 @@
 default_transport = "PikaTransport"
 
 
-def lookup(transport: str) -> type[CommonTransport]:
+def lookup(transport: str | None) -> type[CommonTransport]:
     """Get a transport layer class based on its name."""
-    return get_known_transports().get(
-        transport, get_known_transports()[default_transport]
-    )
+    known_transports = get_known_transports()
+    if transport not in known_transports:
+        return known_transports[default_transport]
+    return known_transports[transport]
 
 
 def add_command_line_options(

src/workflows/transport/common_transport.py

@@ -33,7 +33,7 @@ class CommonTransport:
     #
 
     def __init__(
-        self, middleware: list[type[middleware.BaseTransportMiddleware]] = None
+        self, middleware: list[type[middleware.BaseTransportMiddleware]] | None = None
     ):
         if middleware is None:
             self.middleware = []

src/workflows/transport/middleware/__init__.py

@@ -95,12 +95,12 @@ def transaction_begin(
         return call_next(subscription_id=subscription_id, **kwargs)
 
     def transaction_abort(
-        self, call_next: Callable, transaction_id: int = None, **kwargs
+        self, call_next: Callable, transaction_id: int | None = None, **kwargs
     ):
         call_next(transaction_id, **kwargs)
 
     def transaction_commit(
-        self, call_next: Callable, transaction_id: int = None, **kwargs
+        self, call_next: Callable, transaction_id: int | None = None, **kwargs
     ):
         call_next(transaction_id, **kwargs)
 
@@ -172,7 +172,7 @@ def transaction_commit(self, call_next: Callable, *args, **kwargs):
 
 
 class TimerMiddleware(BaseTransportMiddleware):
-    def __init__(self, logger: logging.Logger = None, level=logging.INFO):
+    def __init__(self, logger: logging.Logger | None = None, level=logging.INFO):
         if logger is None:
             logger = logging.getLogger(__name__)
         self.logger = logger

src/workflows/transport/offline_transport.py

@@ -28,7 +28,7 @@ class OfflineTransport(CommonTransport):
     config: dict[Any, Any] = {}
 
     def __init__(
-        self, middleware: list[type[middleware.BaseTransportMiddleware]] = None
+        self, middleware: list[type[middleware.BaseTransportMiddleware]] | None = None
     ):
         self._connected = False
         super().__init__(middleware=middleware)

src/workflows/transport/pika_transport.py

@@ -63,7 +63,7 @@ class PikaTransport(CommonTransport):
     config: dict[Any, Any] = {}
 
     def __init__(
-        self, middleware: list[type[middleware.BaseTransportMiddleware]] = None
+        self, middleware: list[type[middleware.BaseTransportMiddleware]] | None = None
     ):
         self._channel = None
         self._conn = None

src/workflows/transport/stomp_transport.py

@@ -34,7 +34,7 @@ class StompTransport(CommonTransport):
     config: dict[Any, Any] = {}
 
     def __init__(
-        self, middleware: list[type[middleware.BaseTransportMiddleware]] = None
+        self, middleware: list[type[middleware.BaseTransportMiddleware]] | None = None
     ):
         self._connected = False
         self._namespace = ""