Tecan spark

pylabrobot/io/usb.py

@@ -3,7 +3,7 @@
 import time
 from concurrent.futures import ThreadPoolExecutor
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, List, Optional
+from typing import TYPE_CHECKING, Callable, List, Optional
 
 from pylabrobot.io.capture import Command, capturer, get_capture_or_validation_active
 from pylabrobot.io.errors import ValidationError
@@ -51,6 +51,10 @@ def __init__(
     packet_read_timeout: int = 3,
     read_timeout: int = 30,
     write_timeout: int = 30,
+    configuration_callback: Optional[Callable[["usb.core.Device"], None]] = None,
+    max_workers: int = 1,
+    read_endpoint_address: Optional[int] = None,
+    write_endpoint_address: Optional[int] = None,
   ):
     """Initialize an io.USB object.
 
@@ -63,6 +67,11 @@ def __init__(
       packet_read_timeout: The timeout for reading packets from the machine in seconds.
       read_timeout: The timeout for reading from the machine in seconds.
       write_timeout: The timeout for writing to the machine in seconds.
+      read_endpoint_address: The address of the read endpoint. If `None`, find the first IN endpoint.
+      write_endpoint_address: The address of the write endpoint. If `None`, find the first OUT endpoint.
+      configuration_callback: A callback that takes the device object as an argument and performs
+        any necessary configuration. If `None`, `dev.set_configuration()` is called.
+      max_workers: The maximum number of worker threads for USB I/O operations.
     """
 
     super().__init__()
@@ -82,8 +91,12 @@ def __init__(
     self.packet_read_timeout = packet_read_timeout
     self.read_timeout = read_timeout
     self.write_timeout = write_timeout
+    self.read_endpoint_address = read_endpoint_address
+    self.write_endpoint_address = write_endpoint_address
+    self.configuration_callback = configuration_callback
+    self.max_workers = max_workers
 
-    self.dev: Optional["usb.core.Device"] = None  # TODO: make this a property
+    self.dev: Optional[usb.core.Device] = None  # TODO: make this a property
     self.read_endpoint: Optional[usb.core.Endpoint] = None
     self.write_endpoint: Optional[usb.core.Endpoint] = None
 
@@ -114,13 +127,15 @@ async def write(self, data: bytes, timeout: Optional[float] = None):
       raise RuntimeError("Call setup() first.")
     await loop.run_in_executor(
       self._executor,
-      lambda: dev.write(write_endpoint, data, timeout=timeout),
+      lambda: dev.write(
+        write_endpoint, data, timeout=int(timeout * 1000)
+      ),  # PyUSB expects timeout in milliseconds
     )
     if len(data) % write_endpoint.wMaxPacketSize == 0:
       # send a zero-length packet to indicate the end of the transfer
       await loop.run_in_executor(
         self._executor,
-        lambda: dev.write(write_endpoint, b"", timeout=timeout),
+        lambda: dev.write(write_endpoint, b"", timeout=int(timeout * 1000)),
       )
     logger.log(LOG_LEVEL_IO, "%s write: %s", self._unique_id, data)
     capturer.record(
@@ -131,25 +146,56 @@ async def write(self, data: bytes, timeout: Optional[float] = None):
       )
     )
 
-  def _read_packet(self, size: Optional[int] = None) -> Optional[bytearray]:
+  def _read_packet(
+    self,
+    size: Optional[int] = None,
+    timeout: Optional[float] = None,
+    endpoint: Optional[int] = None,
+  ) -> Optional[bytearray]:
     """Read a packet from the machine.
 
     Args:
       size: The maximum number of bytes to read. If `None`, read up to wMaxPacketSize bytes.
+      timeout: The timeout for reading from the device in seconds. If `None`, use the default
+        timeout (specified by the `packet_read_timeout` attribute).
+      endpoint: The endpoint address to read from. If `None`, use the default read endpoint.
 
     Returns:
       A bytearray containing the data read, or None if no data was received.
     """
 
-    assert self.dev is not None and self.read_endpoint is not None, "Device not connected."
+    assert self.dev is not None, "Device not connected."
 
-    read_size = size if size is not None else self.read_endpoint.wMaxPacketSize
+    ep = endpoint if endpoint is not None else self.read_endpoint
+    if ep is None:
+      raise RuntimeError("Read endpoint not found. Call setup() first.")
+
+    # Get max packet size if size is not provided
+    if size is None:
+      if isinstance(ep, int):
+        # Find endpoint object to get max packet size
+        cfg = self.dev.get_active_configuration()
+        intf = cfg[(0, 0)]
+        ep_obj = usb.util.find_descriptor(
+          intf,
+          custom_match=lambda e: e.bEndpointAddress == ep,
+        )
+        if ep_obj is None:
+          raise ValueError(f"Endpoint 0x{ep:02x} not found.")
+        read_size = ep_obj.wMaxPacketSize
+      else:
+        read_size = ep.wMaxPacketSize
+    else:
+      read_size = size
+
+    if timeout is None:
+      timeout = self.packet_read_timeout
 
     try:
       res = self.dev.read(
-        self.read_endpoint,
+        ep,
         read_size,
-        timeout=int(self.packet_read_timeout * 1000),  # timeout in ms
+        timeout=int(timeout * 1000),  # timeout in ms
       )
 
       if res is not None:
@@ -318,7 +364,7 @@ def ctrl_transfer(
 
     return bytearray(res)
 
-  async def setup(self):
+  async def setup(self, empty_buffer=True):
     """Initialize the USB connection to the machine."""
 
     if self.dev is not None:
@@ -347,22 +393,37 @@ async def setup(self):
 
     # set the active configuration. With no arguments, the first
     # configuration will be the active one
-    self.dev.set_configuration()
+    if self.configuration_callback is not None:
+      self.configuration_callback(self.dev)
+    else:
+      self.dev.set_configuration()
 
     cfg = self.dev.get_active_configuration()
     intf = cfg[(0, 0)]
 
-    self.write_endpoint = usb.util.find_descriptor(
-      intf,
-      custom_match=lambda e: usb.util.endpoint_direction(e.bEndpointAddress)
-      == usb.util.ENDPOINT_OUT,
-    )
+    if self.write_endpoint_address is not None:
+      self.write_endpoint = usb.util.find_descriptor(
+        intf,
+        custom_match=lambda e: e.bEndpointAddress == self.write_endpoint_address,
+      )
+    else:
+      self.write_endpoint = usb.util.find_descriptor(
+        intf,
+        custom_match=lambda e: usb.util.endpoint_direction(e.bEndpointAddress)
+        == usb.util.ENDPOINT_OUT,
+      )
 
-    self.read_endpoint = usb.util.find_descriptor(
-      intf,
-      custom_match=lambda e: usb.util.endpoint_direction(e.bEndpointAddress)
-      == usb.util.ENDPOINT_IN,
-    )
+    if self.read_endpoint_address is not None:
+      self.read_endpoint = usb.util.find_descriptor(
+        intf,
+        custom_match=lambda e: e.bEndpointAddress == self.read_endpoint_address,
+      )
+    else:
+      self.read_endpoint = usb.util.find_descriptor(
+        intf,
+        custom_match=lambda e: usb.util.endpoint_direction(e.bEndpointAddress)
+        == usb.util.ENDPOINT_IN,
+      )
 
     logger.info(
       "Found endpoints. \nWrite:\n %s \nRead:\n %s",
@@ -371,10 +432,11 @@ async def setup(self):
     )
 
     # Empty the read buffer.
-    while self._read_packet() is not None:
-      pass
+    if empty_buffer:
+      while self._read_packet() is not None:
+        pass
 
-    self._executor = ThreadPoolExecutor(max_workers=1)
+    self._executor = ThreadPoolExecutor(max_workers=self.max_workers)
 
   async def stop(self):
     """Close the USB connection to the machine."""
@@ -401,6 +463,8 @@ def serialize(self) -> dict:
       "packet_read_timeout": self.packet_read_timeout,
       "read_timeout": self.read_timeout,
       "write_timeout": self.write_timeout,
+      "read_endpoint_address": self.read_endpoint_address,
+      "write_endpoint_address": self.write_endpoint_address,
     }
 
 
@@ -415,6 +479,10 @@ def __init__(
     packet_read_timeout: int = 3,
     read_timeout: int = 30,
     write_timeout: int = 30,
+    read_endpoint_address: Optional[int] = None,
+    write_endpoint_address: Optional[int] = None,
+    configuration_callback: Optional[Callable[["usb.core.Device"], None]] = None,
+    max_workers: int = 1,
   ):
     super().__init__(
       id_vendor=id_vendor,
@@ -424,10 +492,14 @@ def __init__(
       packet_read_timeout=packet_read_timeout,
       read_timeout=read_timeout,
       write_timeout=write_timeout,
+      read_endpoint_address=read_endpoint_address,
+      write_endpoint_address=write_endpoint_address,
+      configuration_callback=configuration_callback,
+      max_workers=max_workers,
     )
     self.cr = cr
 
-  async def setup(self):
+  async def setup(self, empty_buffer=True):
     pass
 
   async def write(self, data: bytes, timeout: Optional[float] = None):

pylabrobot/plate_reading/__init__.py

@@ -48,3 +48,4 @@
   ImagingResult,
   Objective,
 )
+from .tecan.spark20m.spark_backend import SparkBackend

pylabrobot/plate_reading/tecan/spark20m/controls/__init__.py

@@ -0,0 +1,11 @@
+from .base_control import BaseControl
+from .camera_control import CameraControl
+from .config_control import ConfigControl
+from .data_control import DataControl
+from .injector_control import InjectorControl
+from .measurement_control import MeasurementControl
+from .movement_control import MovementControl
+from .optics_control import OpticsControl
+from .plate_transport_control import PlateControl
+from .sensor_control import SensorControl
+from .system_control import SystemControl

pylabrobot/plate_reading/tecan/spark20m/controls/base_control.py

@@ -0,0 +1,8 @@
+from typing import Awaitable, Callable, Optional
+
+SendCommandFunc = Callable[..., Awaitable[Optional[str]]]
+
+
+class BaseControl:
+  def __init__(self, send_command: SendCommandFunc) -> None:
+    self.send_command = send_command