diff --git a/src/qcodes_loop/data/data_array.py b/src/qcodes_loop/data/data_array.py index 2ea48881..72129728 100644 --- a/src/qcodes_loop/data/data_array.py +++ b/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. diff --git a/src/qcodes_loop/tests/test_data.py b/src/qcodes_loop/tests/test_data.py index c441d296..b6b648ac 100644 --- a/src/qcodes_loop/tests/test_data.py +++ b/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,)