From 614208ebe2ec8d2cee3e2eb6c139172a99bc9dc2 Mon Sep 17 00:00:00 2001 From: "Jens H. Nielsen" Date: Wed, 19 Aug 2026 21:39:11 +0200 Subject: [PATCH 1/2] Annotate DataArray shape, set_arrays, array_id and action_indices These are assigned from unannotated arguments whose defaults are an empty tuple or None, so type checkers inferred the narrowest type that fits the default, such as tuple[()] for set_arrays. Indexing them was therefore an error for consumers. nest is annotated too since it also assigns set_arrays and action_indices, and checkers take the union of every assignment site. shape is declared at its assignment because it is additionally assigned in init_data. Also add an explicit __iter__. Like __len__ it has to be delegated by hand because iter() looks the attribute up on the type rather than going through DelegateAttributes. Iteration already worked via the legacy __getitem__ protocol but the array was not recognised as iterable. --- src/qcodes_loop/data/data_array.py | 34 ++++++++++++++++++++++++------ 1 file changed, 28 insertions(+), 6 deletions(-) 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. From 83f3a8649b3d66e7dba3c6f139984df41187fffb Mon Sep 17 00:00:00 2001 From: "Jens H. Nielsen" Date: Wed, 19 Aug 2026 21:57:08 +0200 Subject: [PATCH 2/2] Test DataArray iteration Iteration was only covered incidentally, by two formatter tests that happen to zip over data arrays in the test body itself. Nothing in qcodes_loop iterates a DataArray, so the only production consumers are downstream, which makes that coverage easy to lose. Assert the semantics directly: 1D yields values, 2D yields rows, and iteration agrees with indexing and length. --- src/qcodes_loop/tests/test_data.py | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) 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,)