Annotate DataArray shape, set_arrays, array_id and action_indices

src/qcodes_loop/data/data_array.py

@@ -4,6 +4,8 @@
 import numpy as np
 
 if TYPE_CHECKING:
+    from collections.abc import Iterator
+
     import xarray as xr
 
 import logging
@@ -122,10 +124,10 @@ def __init__(
         full_name=None,
         label=None,
         snapshot=None,
-        array_id=None,
-        set_arrays=(),
-        shape=None,
-        action_indices=(),
+        array_id: "str | None" = None,
+        set_arrays: "tuple[DataArray, ...]" = (),
+        shape: "tuple[int, ...] | None" = None,
+        action_indices: "tuple[int, ...]" = (),
         unit=None,
         units=None,
         is_setpoint=False,
@@ -134,7 +136,11 @@ def __init__(
         self.name = name
         self.full_name = full_name or name
         self.label = label
-        self.shape = shape
+        # Declared explicitly because ``nest`` and ``init_data`` also assign to
+        # it; without a declaration type checkers take the union of every
+        # assignment, which includes fixed length tuples such as ``tuple[()]``
+        # and makes indexing the shape an error for consumers.
+        self.shape: tuple[int, ...] | None = shape
         if units is not None:
             _LOG.warning(
                 f"`units` is deprecated for the "
@@ -210,7 +216,12 @@ def data_set(self, new_data_set):
             raise RuntimeError("A DataArray can only be part of one DataSet")
         self._data_set = new_data_set
 
-    def nest(self, size, action_index=None, set_array=None):
+    def nest(
+        self,
+        size: int,
+        action_index: "int | None" = None,
+        set_array: "DataArray | None" = None,
+    ):
         """
         Nest this array inside a new outer loop.
 
@@ -381,6 +392,17 @@ def __len__(self):
         """
         return len(self.ndarray)
 
+    def __iter__(self) -> "Iterator[Any]":
+        """
+        Iterate over the values in this array.
+
+        Must be explicitly delegated, because iter() looks up ``__iter__`` on
+        the type rather than the instance. Without it iteration still works via
+        the legacy ``__getitem__`` protocol, but the array is not recognised as
+        iterable by type checkers.
+        """
+        return iter(self.ndarray)
+
     def flat_index(self, indices, index_fill=None):
         """
         Generate the raveled index for the given indices.

src/qcodes_loop/tests/test_data.py

@@ -134,6 +134,28 @@ def test_preset_data(self):
         self.assertEqual(data2.ndarray.tolist(), list2d)
         self.assertEqual(data2.shape, (2, 2))
 
+    def test_iteration(self):
+        # a 1D array iterates over its individual values
+        data = DataArray(preset_data=[1.0, 2.0, 3.0])
+        self.assertEqual([float(value) for value in data], [1.0, 2.0, 3.0])
+
+        # iteration agrees with indexing and length
+        self.assertEqual(len(list(data)), len(data))
+        self.assertEqual(
+            [float(data[i]) for i in range(len(data))],
+            [float(value) for value in data],
+        )
+
+        # a 2D array iterates over its rows, matching numpy. Consumers such as
+        # the qcodes legacy dataset importer rely on this to loop over the
+        # outer setpoints of a 2D array.
+        data2d = DataArray(preset_data=[[1.0, 2.0], [3.0, 4.0]])
+        self.assertEqual([row.tolist() for row in data2d], [[1.0, 2.0], [3.0, 4.0]])
+
+        # an array that has no data yet cannot be iterated
+        with self.assertRaises(TypeError):
+            iter(DataArray(name="no_data"))
+
     def test_init_data_error(self):
         data = DataArray(preset_data=[1, 2])
         data.shape = (3,)