error handling: option 1 (PoC)

pylabrobot/error_handling/__init__.py

@@ -0,0 +1,2 @@
+from .handles_errors import handles_errors
+from .serial_handler import SerialErrorHandler

pylabrobot/error_handling/handles_errors.py

@@ -0,0 +1,26 @@
+import functools
+import inspect
+
+def handles_errors(func):
+  @functools.wraps(func)
+  async def wrapper(self, *args, **kwargs):
+    try:
+      return await func(self, *args, **kwargs)
+    except Exception as error:
+      handler = self._error_handlers.get(type(error))
+      if handler:
+        print(f"Handling error {error} with: {handler}")
+        # bind the wrapper to this instance so that
+        # retries still go through the decorator
+        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"}
+
+        # call the handler, passing it the *decorated* method
+        return await handler(bound, error, **bound_args)
+      # no handler registered -> re‑raise
+      raise
+  return wrapper

pylabrobot/error_handling/serial_handler.py

@@ -0,0 +1,13 @@
+class SerialErrorHandler:
+  def __init__(self, child_handlers: list):
+    self.child_handlers = child_handlers
+    self.fallback = fallback
+    self.index = 0
+
+  def __call__(self, func, *args, **kwargs):
+    print("serial error handler is choosing next child handler")
+    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 handler(func, *args, **kwargs)

pylabrobot/liquid_handling/error_handlers.py

@@ -0,0 +1,23 @@
+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

@@ -19,6 +19,7 @@
   Sequence,
   Set,
   Tuple,
+  Type,
   Union,
   cast,
 )
@@ -60,6 +61,7 @@
 from pylabrobot.resources.liquid import Liquid
 from pylabrobot.resources.rotation import Rotation
 from pylabrobot.tilting.tilter import Tilter
+from pylabrobot.error_handling import handles_errors
 
 from .backends import LiquidHandlerBackend
 from .standard import (
@@ -147,6 +149,8 @@ def __init__(self, backend: LiquidHandlerBackend, deck: Deck):
 
     self._resource_pickup: Optional[ResourcePickup] = None
 
+    self._error_handlers: Dict[Type[Exception], Callable] = {}
+
   async def setup(self, **backend_kwargs):
     """Prepare the robot for use."""
 
@@ -333,7 +337,8 @@ 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
+  @handles_errors
+  # @need_setup_finished
   async def pick_up_tips(
     self,
     tip_spots: List[TipSpot],
@@ -1180,6 +1185,9 @@ async def transfer(
         **backend_kwargs,
       )
 
+    if error is not None:
+      raise error
+
   @contextlib.contextmanager
   def use_channels(self, channels: List[int]):
     """Temporarily use the specified channels as a default argument to `use_channels`.
@@ -2146,6 +2154,21 @@ async def move_channel_z(self, channel: int, z: float):
     """Move channel to absolute z position"""
     assert 0 <= channel < self.backend.num_channels, f"Invalid channel: {channel}"
     await self.backend.move_channel_z(channel=channel, z=z)
+
+  @contextlib.contextmanager
+  def on_fail(self, error_cls: Type[Exception], handler: Callable):
+    """Register a handler to be called when an error occurs.
+
+    Args:
+      error_cls: The exception class to handle.
+      handler: The handler function to call.
+    """
+
+    self._error_handlers[error_cls] = handler
+    try:
+      yield
+    finally:
+      del self._error_handlers[error_cls]
 
   # -- Resource methods --