Don't try to index a flat array with two indices (#475)

* Don't try to index a flat array with two indices * Add a test to check that flat dataset repr doesn't bork * Check that the repr is outputting the right length of flat array * Changelog --------- Co-authored-by: Stuart Mumford <[email protected]>
DKISTDC · Jan 6, 2025 · cd50651 · cd50651
1 parent 66035ba
commit cd50651
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 1 deletion.
diff --git a/changelog/475.trivial.rst b/changelog/475.trivial.rst
@@ -0,0 +1 @@
+Fix small bug which caused `ds.flat` to break if not indexed.
diff --git a/dkist/dataset/tests/test_dataset.py b/dkist/dataset/tests/test_dataset.py
@@ -51,6 +51,12 @@ def test_repr(dataset, dataset_3d):
     assert str(dataset_3d.data) in r
 
 
+@pytest.mark.accept_cli_dataset
+def test_flat_repr(large_tiled_dataset):
+    r = repr(large_tiled_dataset.flat)
+    assert f"is an array of ({np.prod(large_tiled_dataset.shape)},) Dataset objects" in r
+
+
 @pytest.mark.accept_cli_dataset
 def test_wcs_roundtrip(dataset):
     p = [1*u.pixel] * dataset.wcs.pixel_n_dim

diff --git a/dkist/dataset/utils.py b/dkist/dataset/utils.py
@@ -18,7 +18,7 @@ def dataset_info_str(ds_in):
     dstype = type(ds_in).__name__
     if is_tiled:
         tile_shape = ds_in.shape
-        ds = ds_in[0, 0]
+        ds = ds_in.flat[0]
     else:
         ds = ds_in
     wcs = ds.wcs.low_level_wcs
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1 @@
		Fix small bug which caused `ds.flat` to break if not indexed.