Type recipe package

Author: ndevenish

src/workflows/recipe/__init__.py

@@ -26,15 +26,15 @@
 
 def _wrap_subscription(
     transport_layer: CommonTransport,
-    subscription_call,
-    channel,
-    callback,
-    *args,
+    subscription_call: Callable[..., int],
+    channel: str,
+    callback: Callable[..., Any],
+    *args: Any,
     mangle_for_receiving: Callable[[dict[str, Any]], dict[str, Any]] | None = None,
     allow_non_recipe_messages: bool = False,
-    log_extender=None,
-    **kwargs,
-):
+    log_extender: Callable[[str, Any], AbstractContextManager[Any]] | None = None,
+    **kwargs: Any,
+) -> int:
     """Internal method to create an intercepting function for incoming messages
     to interpret recipes. This function is then used to subscribe to a channel
     on the transport layer.
@@ -61,7 +61,7 @@ def _wrap_subscription(
     """
 
     @functools.wraps(callback)
-    def unwrap_recipe(header, message: dict[str, Any]):
+    def unwrap_recipe(header: dict[str, Any], message: dict[str, Any]) -> Any:
         """Unpack incoming messages when they are in a recipe format.
 
         Other messages are passed through unmodified.
@@ -124,11 +124,11 @@ def wrap_subscribe(
     transport_layer: CommonTransport,
     channel: str,
     callback: Callable[[RecipeWrapper, dict, dict], None],
-    *args,
+    *args: Any,
     allow_non_recipe_messages: Literal[False] = False,
     mangle_for_receiving: Callable[[Any], Any] | None = None,
     log_extender: Callable[[str, Any], AbstractContextManager[Any]] | None = None,
-    **kwargs,
+    **kwargs: Any,
 ) -> int: ...
 
 
@@ -137,24 +137,24 @@ def wrap_subscribe(
     transport_layer: CommonTransport,
     channel: str,
     callback: Callable[[RecipeWrapper | None, dict, dict | bytes], None],
-    *args,
+    *args: Any,
     allow_non_recipe_messages: Literal[True],
     mangle_for_receiving: Callable[[Any], Any] | None = None,
     log_extender: Callable[[str, Any], AbstractContextManager[Any]] | None = None,
-    **kwargs,
+    **kwargs: Any,
 ) -> int: ...
 
 
 def wrap_subscribe(
-    transport_layer,
-    channel,
-    callback,
-    *args,
-    allow_non_recipe_messages=False,
-    mangle_for_receiving=None,
-    log_extender=None,
-    **kwargs,
-):
+    transport_layer: CommonTransport,
+    channel: str,
+    callback: Callable[..., Any],
+    *args: Any,
+    allow_non_recipe_messages: bool = False,
+    mangle_for_receiving: Callable[[Any], Any] | None = None,
+    log_extender: Callable[[str, Any], AbstractContextManager[Any]] | None = None,
+    **kwargs: Any,
+) -> int:
     """Listen to a queue on the transport layer, similar to the subscribe call in
     transport/common_transport.py. Intercept all incoming messages and parse
     for recipe information.
@@ -186,13 +186,13 @@ def wrap_subscribe(
 
 
 def wrap_subscribe_broadcast(
-    transport_layer,
-    channel,
-    callback,
-    *args,
+    transport_layer: CommonTransport,
+    channel: str,
+    callback: Callable[..., Any],
+    *args: Any,
     mangle_for_receiving: Callable[[Any], Any] | None = None,
-    **kwargs,
-):
+    **kwargs: Any,
+) -> int:
     """Listen to a topic on the transport layer, similar to the
     subscribe_broadcast call in transport/common_transport.py. Intercept all
     incoming messages and parse for recipe information.

src/workflows/recipe/recipe.py

@@ -132,7 +132,7 @@ def validate(self) -> None:
         # Detect cycles
         touched_nodes = {"start", "error"}
 
-        def flatten_links(struct):
+        def flatten_links(struct: Any) -> list[int]:
             """Take an output/error link object, list or dictionary and return flat list of linked nodes."""
             if struct is None:
                 return []
@@ -153,7 +153,7 @@ def flatten_links(struct):
                 "Invalid recipe: Invalid link in recipe (%s)" % str(struct)
             )
 
-        def find_cycles(path):
+        def find_cycles(path: list[Any]) -> None:
             """Depth-First-Search helper function to identify cycles."""
             if path[-1] not in self.recipe:
                 raise workflows.Error(
@@ -215,23 +215,23 @@ def apply_parameters(self, parameters: dict[str, Any]) -> None:
         """
 
         class SafeString:
-            def __init__(self, s):
+            def __init__(self, s: str):
                 self.string = s
 
-            def __repr__(self):
+            def __repr__(self) -> str:
                 return "{" + self.string + "}"
 
-            def __str__(self):
+            def __str__(self) -> str:
                 return "{" + self.string + "}"
 
-            def __getitem__(self, item):
+            def __getitem__(self, item: str) -> SafeString:
                 return SafeString(self.string + "[" + item + "]")
 
         class SafeDict(dict):
             """A dictionary that returns undefined keys as {keyname}.
             This can be used to selectively replace variables in datastructures."""
 
-            def __missing__(self, key):
+            def __missing__(self, key: str) -> SafeString:
                 return SafeString(key)
 
         # By default the python formatter class is used to resolve {item} references
@@ -242,15 +242,15 @@ def __missing__(self, key):
         # string.
         ds_formatter = string.Formatter()
 
-        def ds_format_field(value, spec):
+        def ds_format_field(value: Any, spec: str) -> str:
             ds_format_field.last = value  # type: ignore
             return ""
 
         ds_formatter.format_field = ds_format_field  # type: ignore
 
         params = SafeDict(parameters)
 
-        def _recursive_apply(item):
+        def _recursive_apply(item: Any) -> Any:
             """Helper function to recursively apply replacements."""
             if isinstance(item, str):
                 if item.startswith("{$REPLACE") and item.endswith("}"):
@@ -323,7 +323,7 @@ def merge(self, other: Recipe | str) -> Recipe:
                 new_recipe[translation[key]] = value
 
         # Rewrite all copied entries to point to new keys
-        def translate(x):
+        def translate(x: Any) -> Any:
             if isinstance(x, list):
                 return list(map(translate, x))
             elif isinstance(x, tuple):

src/workflows/recipe/validate.py

@@ -54,7 +54,7 @@ def validate_recipe(json_filename: str | os.PathLike[str]) -> None:
         raise e
 
 
-def main():
+def main() -> None:
     """Run the program from entry point"""
     parser = argparse.ArgumentParser()
     parser.add_argument(

src/workflows/recipe/wrapper.py

@@ -301,7 +301,7 @@ def start(
         header: dict[str, Any] | None = None,
         *,
         mangle_for_sending: Callable[[Any], Any] | None = None,
-        **kwargs,
+        **kwargs: Any,
     ) -> None:
         """
         Trigger the start of a recipe.
@@ -375,7 +375,7 @@ def checkpoint(
                 usually "encode to JSON".
             **kwargs: Keywords passed on to the transport.
         """
-        if not self.recipe_step:
+        if not self.recipe_step or self.recipe_pointer is None:
             raise ValueError(
                 "This RecipeWrapper object does not contain "
                 "a recipe with a selected step."
@@ -418,7 +418,9 @@ def apply_parameters(self, parameters: dict[str, Any]) -> None:
         assert self.recipe_pointer is not None
         self.recipe_step = self.recipe[self.recipe_pointer]
 
-    def _generate_full_recipe_message(self, destination, message, add_path_step):
+    def _generate_full_recipe_message(
+        self, destination: int, message: Any, add_path_step: bool
+    ) -> dict[str, Any]:
         """Factory function to generate independent message objects for
         downstream recipients with different destinations."""
         if add_path_step and self.recipe_pointer:
@@ -436,17 +438,17 @@ def _generate_full_recipe_message(self, destination, message, add_path_step):
 
     def _send_to_destinations(
         self,
-        destinations,
-        message,
-        header=None,
+        destinations: int | list[int],
+        message: Any,
+        header: dict[str, Any] | None = None,
         mangle_for_sending: Callable[[Any], Any] | None = None,
-        **kwargs,
-    ):
+        **kwargs: Any,
+    ) -> None:
         """Send messages to a list of numbered destinations. This is an internal
         helper method used by the public 'send' methods.
         """
         if not isinstance(destinations, list):
-            destinations = (destinations,)
+            destinations = [destinations]
         for destination in destinations:
             self._send_to_destination(
                 destination,
@@ -458,13 +460,13 @@ def _send_to_destinations(
 
     def _send_to_destination(
         self,
-        destination,
-        header,
-        payload,
-        transport_kwargs,
-        add_path_step=True,
+        destination: int,
+        header: dict[str, Any] | None,
+        payload: Any,
+        transport_kwargs: dict[str, Any],
+        add_path_step: bool = True,
         mangle_for_sending: Callable[[Any], Any] | None = None,
-    ):
+    ) -> None:
         """Helper function to send a message to a specific recipe destination."""
         if header:
             header = header.copy()
@@ -516,7 +518,9 @@ def _send_to_destination(
                 **dest_kwargs,
             )
 
-    def _retry_transport(self, function, *args, **kwargs):
+    def _retry_transport(
+        self, function: Callable[..., Any], *args: Any, **kwargs: Any
+    ) -> Any:
         """Attempt to send a message, and in case the connection has been lost,
         attempt to reconnect. Reconnecting only works on the assumption that
         the previous connection did not include any subscriptions, which should