linear_tip_spot_generator is a class (#395)

Author: rickwierenga

docs/user_guide/tip-spot-generators.ipynb

@@ -37,7 +37,7 @@
         {
           "data": {
             "text/plain": [
-              "TipSpot(name=tip_rack_0_tipspot_0_0, location=(007.200, 068.300, -83.500), size_x=9.0, size_y=9.0, size_z=0, category=tip_spot)"
+              "TipSpot(name=tip_rack_0_tipspot_0_0, location=Coordinate(007.200, 068.300, -83.500), size_x=9.0, size_y=9.0, size_z=0, category=tip_spot)"
             ]
           },
           "execution_count": 2,
@@ -88,7 +88,7 @@
         {
           "data": {
             "text/plain": [
-              "TipSpot(name=tip_rack_0_tipspot_0_0, location=(007.200, 068.300, -83.500), size_x=9.0, size_y=9.0, size_z=0, category=tip_spot)"
+              "TipSpot(name=tip_rack_0_tipspot_0_0, location=Coordinate(007.200, 068.300, -83.500), size_x=9.0, size_y=9.0, size_z=0, category=tip_spot)"
             ]
           },
           "execution_count": 4,
@@ -129,6 +129,50 @@
         "[ts.name for ts in tip_spots]"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "Save the state of the generator at an arbitrary point by calling `save_state`. This method will be called automatically when the program crashes or is interrupted."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 6,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "linear_generator.save_state()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "Override the index by calling `set_index`."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 7,
+      "metadata": {},
+      "outputs": [
+        {
+          "data": {
+            "text/plain": [
+              "TipSpot(name=tip_rack_0_tipspot_1_4, location=Coordinate(016.200, 032.300, -83.500), size_x=9.0, size_y=9.0, size_z=0, category=tip_spot)"
+            ]
+          },
+          "execution_count": 7,
+          "metadata": {},
+          "output_type": "execute_result"
+        }
+      ],
+      "source": [
+        "linear_generator.set_index(12)\n",
+        "await linear_generator.__anext__()"
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {},
@@ -140,7 +184,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 6,
+      "execution_count": 8,
       "metadata": {},
       "outputs": [],
       "source": [
@@ -153,16 +197,16 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 7,
+      "execution_count": 9,
       "metadata": {},
       "outputs": [
         {
           "data": {
             "text/plain": [
-              "TipSpot(name=tip_rack_0_tipspot_0_2, location=(007.200, 050.300, -83.500), size_x=9.0, size_y=9.0, size_z=0, category=tip_spot)"
+              "TipSpot(name=tip_rack_0_tipspot_0_3, location=Coordinate(007.200, 041.300, -83.500), size_x=9.0, size_y=9.0, size_z=0, category=tip_spot)"
             ]
           },
-          "execution_count": 7,
+          "execution_count": 9,
           "metadata": {},
           "output_type": "execute_result"
         }
@@ -188,7 +232,7 @@
       "name": "python",
       "nbconvert_exporter": "python",
       "pygments_lexer": "ipython3",
-      "version": "3.9.19"
+      "version": "3.13.1"
     }
   },
   "nbformat": 4,

pylabrobot/resources/functional.py

@@ -1,3 +1,4 @@
+import atexit
 import json
 import logging
 import os
@@ -14,31 +15,52 @@ def get_all_tip_spots(tip_racks: List[TipRack]) -> List[TipSpot]:
   return [spot for rack in tip_racks for spot in rack.get_all_items()]
 
 
-async def linear_tip_spot_generator(
-  tip_spots: List[TipSpot],
-  cache_file_path: Optional[str] = None,
-  repeat: bool = False,
-) -> AsyncGenerator[TipSpot, None]:
-  """Tip spot generator with disk caching. Linearly iterate through all tip spots and
-  raise StopIteration when all spots have been used."""
-  tip_spot_idx = 0
-  if cache_file_path is not None and os.path.exists(cache_file_path):
-    with open(cache_file_path, "r", encoding="utf-8") as f:
-      data = json.load(f)
-      tip_spot_idx = data["tip_spot_idx"]
-      logger.info("loaded tip idx from disk: %s", data)
-
-  while True:
-    if cache_file_path is not None:
-      with open(cache_file_path, "w", encoding="utf-8") as f:
-        json.dump({"tip_spot_idx": tip_spot_idx}, f)
-    yield tip_spots[tip_spot_idx]
-    tip_spot_idx += 1
-    if tip_spot_idx >= len(tip_spots):
-      if repeat:
-        tip_spot_idx = 0
-      else:
-        return
+class linear_tip_spot_generator:
+  def __init__(
+    self, tip_spots: List[TipSpot], cache_file_path: Optional[str] = None, repeat: bool = False
+  ):
+    self.tip_spots = tip_spots
+    self.cache_file_path = cache_file_path
+    self.repeat = repeat
+    self._tip_spot_idx = 0
+
+    if self.cache_file_path and os.path.exists(self.cache_file_path):
+      try:
+        with open(self.cache_file_path, "r", encoding="utf-8") as f:
+          data = json.load(f)
+          self._tip_spot_idx = data.get("tip_spot_idx", 0)
+          logger.info("Loaded tip idx from disk: %s", data)
+      except Exception as e:
+        logger.error("Failed to load cache file: %s", e)
+
+    atexit.register(self.save_state)
+
+  def __aiter__(self):
+    return self
+
+  async def __anext__(self) -> TipSpot:
+    while True:
+      self.save_state()
+      if self._tip_spot_idx >= len(self.tip_spots):
+        if self.repeat:
+          self._tip_spot_idx = 0
+        else:
+          raise StopAsyncIteration
+
+      self._tip_spot_idx += 1
+      return self.tip_spots[self._tip_spot_idx - 1]
+
+  def save_state(self):
+    if self.cache_file_path:
+      try:
+        with open(self.cache_file_path, "w", encoding="utf-8") as f:
+          json.dump({"tip_spot_idx": self._tip_spot_idx}, f)
+          logger.info("Saved tip idx to disk: %s", self._tip_spot_idx)
+      except Exception as e:
+        logger.error("Failed to save cache file: %s", e)
+
+  def set_index(self, index: int):
+    self._tip_spot_idx = index
 
 
 async def randomized_tip_spot_generator(