added iterators to the list

src/ophyd_async/core/_detector.py

@@ -3,7 +3,7 @@
 import asyncio
 import time
 from abc import ABC, abstractmethod
-from collections.abc import AsyncGenerator, AsyncIterator, Callable, Sequence
+from collections.abc import AsyncGenerator, AsyncIterator, Callable, Iterator, Sequence
 from enum import Enum
 from typing import (
     Generic,
@@ -206,8 +206,9 @@ def __init__(
         self._watchers: list[Callable] = []
         self._fly_status: WatchableAsyncStatus | None = None
         self._fly_start: float
+        self._frames_to_complete: int = 0
         self._frames_collected: int = 0
-        self._current_number_index: int = 0
+        self._number_of_triggers_iter: Iterator[int]
         self._initial_frame: int
         self._last_frame: int
         self._total_frames_to_capture: int
@@ -317,6 +318,11 @@ async def prepare(self, value: TriggerInfo) -> None:
                 f"but trigger logic provides only {value.deadtime}s"
             )
         self._trigger_info = value
+        self._number_of_triggers_iter = (
+            iter(self._trigger_info.number_of_triggers)
+            if isinstance(self._trigger_info.number_of_triggers, list)
+            else iter([self._trigger_info.number_of_triggers])
+        )
         self._initial_frame = await self.writer.get_indices_written()
         if isinstance(self._trigger_info.number_of_triggers, list):
             self._total_frames_to_capture = sum(self._trigger_info.number_of_triggers)
@@ -341,12 +347,9 @@ async def kickoff(self):
                 f"Kickoff called more than the configured number of "
                 f"{self._total_frames_to_capture} iteration(s)!"
             )
-        if isinstance(self._trigger_info.number_of_triggers, list):
-            self._frames_collected += self._trigger_info.number_of_triggers[
-                self._current_number_index
-            ]
-        else:
-            self._frames_collected += self._trigger_info.number_of_triggers
+
+        self._frames_to_complete = next(self._number_of_triggers_iter)
+        self._frames_collected += self._frames_to_complete
 
     @WatchableAsyncStatus.wrap
     async def complete(self):
@@ -360,7 +363,7 @@ async def complete(self):
             )
         ):
             target = (
-                self._trigger_info.number_of_triggers[self._current_number_index]
+                self._frames_to_complete
                 if isinstance(self._trigger_info.number_of_triggers, list)
                 else self._trigger_info.number_of_triggers
             )
@@ -374,11 +377,7 @@ async def complete(self):
                 time_elapsed=time.monotonic() - self._fly_start,
             )
             if isinstance(self._trigger_info.number_of_triggers, list):
-                if (
-                    index
-                    >= self._trigger_info.number_of_triggers[self._current_number_index]
-                ):
-                    self._current_number_index += 1
+                if index >= self._frames_to_complete:
                     break
             else:
                 if index >= self._trigger_info.number_of_triggers: