error handling: option 3 (PoC)

pylabrobot/error_handling/__init__.py

@@ -0,0 +1,2 @@
+from .handlers import choose_handler, serial_error_handler, until_success, basic_retry_handler
+from .with_error_handler import with_error_handler

pylabrobot/error_handling/handlers/__init__.py

@@ -0,0 +1,4 @@
+from .choose_handler import choose_handler
+from .serial_handler import serial_error_handler
+from .until_success import until_success
+from .retry import basic_retry_handler

pylabrobot/error_handling/handlers/choose_handler.py

@@ -0,0 +1,12 @@
+from typing import Callable, Dict
+
+
+def choose_handler(handlers: Dict[Exception, Callable]) -> Callable:
+  """Choose the appropriate error handler based on the type of error."""
+
+  async def handler(func, exception, **kwargs):
+    for exc_type, handler in handlers.items():
+      if isinstance(exception, exc_type):
+        return await handler(func, exception, **kwargs)
+
+  return handler

pylabrobot/error_handling/handlers/retry.py

@@ -0,0 +1,3 @@
+async def basic_retry_handler(func, error, **kwargs):
+  """Will simply retry the function call with the same arguments."""
+  return await func(**kwargs) 

pylabrobot/error_handling/handlers/serial_handler.py

@@ -0,0 +1,11 @@
+class serial_error_handler:
+  def __init__(self, child_handlers: list):
+    self.child_handlers = child_handlers
+    self.index = 0
+
+  async def __call__(self, func, exception, **kwargs):
+    if self.index >= len(self.child_handlers):
+      raise RuntimeError("No more child handlers to call")
+    handler = self.child_handlers[self.index]
+    self.index += 1
+    return await handler(func, exception, **kwargs)

pylabrobot/error_handling/handlers/until_success.py

@@ -0,0 +1,24 @@
+from typing import Callable
+from typing import Callable, Optional
+
+
+class until_success:
+  """
+  Error handler that retries the given handler until the main function does not raise
+  an exception, or until the maximum number of tries is reached.
+
+  Args:
+    handler: The async function to be executed.
+    max_tries: Maximum number of retries. Default is None, which means infinite retries.
+  """
+
+  def __init__(self, handler: Callable, max_tries: Optional[int] = None):
+    self.handler = handler
+    self.max_tries = max_tries
+    self.attempts = 0
+
+  async def __call__(self, *args, **kwargs):
+    if self.max_tries is not None and self.attempts >= self.max_tries:
+      raise RuntimeError("Maximum number of retries reached")
+    self.attempts += 1
+    return await self.handler(*args, **kwargs)

pylabrobot/error_handling/with_error_handler.py

@@ -0,0 +1,32 @@
+from __future__ import annotations
+
+import functools
+import inspect
+from typing import Any, Awaitable, Callable, Optional, ParamSpec, TypeVar
+
+_P = ParamSpec("_P")
+_R = TypeVar("_R", bound=Awaitable[Any])
+
+Handler: Callable[[Callable[_P, _R], Exception, dict[str, Any]], Awaitable[Any]]
+
+
+def with_error_handler(func: Callable[_P, _R]) -> Callable[_P, _R]:
+  @functools.wraps(func)
+  async def wrapper(self, *args, error_handler: Optional[Handler] = None, **kwargs):
+    try:
+      return await func(self, *args, **kwargs)
+    except Exception as error:
+      print("caught error", error)
+      if error_handler is not None:
+        bound = wrapper.__get__(self, type(self))
+
+        # convert all args to kwargs, remove self
+        sig = inspect.signature(func)
+        bound_args = sig.bind(self, *args, **kwargs)
+        bound_args = {k: v for k, v in bound_args.arguments.items() if k != "self"}
+        bound_args["error_handler"] = error_handler
+
+        return await error_handler(bound, error, **bound_args)
+      raise
+
+  return wrapper

pylabrobot/liquid_handling/error_handlers.py

@@ -0,0 +1,24 @@
+from pylabrobot.liquid_handling.errors import ChannelizedError
+
+
+def try_next_tip_spot(try_tip_spots):
+  async def handler(func, error: Exception, **kwargs):
+    assert isinstance(error, ChannelizedError)
+
+    new_tip_spots, new_use_channels = [], []
+
+    tip_spots = kwargs.pop("tip_spots")
+    if "use_channels" not in kwargs:
+      use_channels = list(range(len(tip_spots)))
+    else:
+      use_channels = kwargs.pop("use_channels")
+
+    for idx, channel_idx in zip(tip_spots, use_channels):
+      if channel_idx in error.errors.keys():
+        new_tip_spots.append(next(try_tip_spots))
+        new_use_channels.append(channel_idx)
+
+    print(f"Retrying with tip spots: {new_tip_spots} and use channels: {new_use_channels}")
+    return await func(tip_spots=new_tip_spots, use_channels=new_use_channels, **kwargs)
+
+  return handler

pylabrobot/liquid_handling/liquid_handler.py

@@ -23,6 +23,7 @@
   cast,
 )
 
+from pylabrobot.error_handling import with_error_handler
 from pylabrobot.liquid_handling.errors import ChannelizedError
 from pylabrobot.liquid_handling.strictness import (
   Strictness,
@@ -333,7 +334,7 @@ def _make_sure_channels_exist(self, channels: List[int]):
     if not len(invalid_channels) == 0:
       raise ValueError(f"Invalid channels: {invalid_channels}")
 
-  @need_setup_finished
+  @with_error_handler
   async def pick_up_tips(
     self,
     tip_spots: List[TipSpot],