From 63281dff735dbf4e23cf7f55c50bd43a94a0e9cb Mon Sep 17 00:00:00 2001
From: Thomas Robitaille <thomas.robitaille@gmail.com>
Date: Fri, 12 Jul 2024 14:03:16 +0100
Subject: [PATCH] Added a test (originally in glue-qt) for IndexedData in the
 image viewer test suite, and fix bug

---
 glue/viewers/image/state.py             |  4 ++--
 glue/viewers/image/tests/test_viewer.py | 30 +++++++++++++++++++++++++
 2 files changed, 32 insertions(+), 2 deletions(-)

diff --git a/glue/viewers/image/state.py b/glue/viewers/image/state.py
index 3a5c6db8b..194650c93 100644
--- a/glue/viewers/image/state.py
+++ b/glue/viewers/image/state.py
@@ -1,7 +1,7 @@
 import uuid
 from collections import defaultdict
 
-from glue.core import BaseData
+from glue.core import BaseData, Data
 from glue.config import colormaps
 from glue.viewers.matplotlib.state import (MatplotlibDataViewerState,
                                            MatplotlibLayerState,
@@ -620,7 +620,7 @@ def reset_contrast_bias(self):
 
     def _update_attribute_display_unit_choices(self, *args):
 
-        if self.layer is None or self.attribute is None:
+        if self.layer is None or self.attribute is None or not isinstance(self.layer, Data):
             ImageLayerState.attribute_display_unit.set_choices(self, [])
             return
 
diff --git a/glue/viewers/image/tests/test_viewer.py b/glue/viewers/image/tests/test_viewer.py
index bf4738de9..7360814d6 100644
--- a/glue/viewers/image/tests/test_viewer.py
+++ b/glue/viewers/image/tests/test_viewer.py
@@ -6,6 +6,7 @@
 from glue.core.data import Data
 from glue.core.link_helpers import LinkSame
 from glue.core.data_region import RegionData
+from glue.core.data_derived import IndexedData
 from astropy.wcs import WCS
 
 from shapely.geometry import Polygon, MultiPolygon, Point
@@ -265,3 +266,32 @@ def test_flipped_wcs_viewer(self):
         assert np.array_equal(original_path_patch, np.flip(new_path_patch, axis=1))
 
         return self.viewer.figure
+
+
+def test_indexed_data():
+
+    # Make sure that the image viewer works properly with IndexedData objects
+
+    data_4d = Data(label='hypercube_wcs',
+                   x=np.random.random((3, 5, 4, 3)),
+                   coords=WCS(naxis=4))
+
+    data_2d = IndexedData(data_4d, (2, None, 3, None))
+
+    application = Application()
+
+    session = application.session
+
+    hub = session.hub
+
+    data_collection = session.data_collection
+    data_collection.append(data_4d)
+    data_collection.append(data_2d)
+
+    viewer = application.new_data_viewer(SimpleImageViewer)
+    viewer.add_data(data_2d)
+
+    assert viewer.state.x_att is data_2d.pixel_component_ids[1]
+    assert viewer.state.y_att is data_2d.pixel_component_ids[0]
+    assert viewer.state.x_att_world is data_2d.world_component_ids[1]
+    assert viewer.state.y_att_world is data_2d.world_component_ids[0]