Make add_data robust to child data

Author: MatthewPortman

jdaviz/configs/default/plugins/data_menu/data_menu.py

@@ -691,10 +691,17 @@ def add_data(self, *data_labels):
         *data_labels : str
             The label(s) of the dataset to add to the viewer.
         """
+        available = self.dataset.choices
+
+        for label in data_labels:
+            parent_label = self.app._get_assoc_data_parent(label)
+            available += self.app._get_assoc_data_children(parent_label)
+
         unavailable = [data_label for data_label in data_labels
-                       if data_label not in self.dataset.choices]
+                       if data_label not in available]
         if len(unavailable):
-            raise ValueError(f"Data labels {unavailable} not able to be loaded into '{self.viewer_id}'.  Must be one of: {self.dataset.choices}")  # noqa
+            raise ValueError(f"Data labels {unavailable} not able to be loaded into '{self.viewer_id}'. Must be one of: {available}")  # noqa
+
         for data_label in data_labels:
             self._app.add_data_to_viewer(self.viewer_id, data_label)
 

jdaviz/core/loaders/tests/test_loaders.py

@@ -813,6 +813,7 @@ def test_load_image_parent_default_auto(self, deconfigged_helper):
         ldr.load()
 
         assert len(deconfigged_helper._app.data_collection) == 4
+        assert len(deconfigged_helper.viewers['Image'].data_menu.data_labels_loaded) == 4
         assert deconfigged_helper._app._get_assoc_data_parent('Image[ERR,1]') == 'Image[SCI,1]'
         assert deconfigged_helper._app._get_assoc_data_children('Image[SCI,1]') == ['Image[ERR,1]']
 
@@ -828,6 +829,7 @@ def test_load_image_parent_none(self, deconfigged_helper):
         ldr.load()
 
         assert len(deconfigged_helper._app.data_collection) == 4
+        assert len(deconfigged_helper.viewers['Image'].data_menu.data_labels_loaded) == 4
         for label in ('Image[SCI,1]', 'Image[ERR,1]', 'Image[SCI,2]', 'Image[ERR,2]'):
             assert deconfigged_helper._app._get_assoc_data_parent(label) is None
             assert deconfigged_helper._app._get_assoc_data_children(label) == []
@@ -859,9 +861,14 @@ def test_load_image_auto_parent_existing_in_dc(self, deconfigged_helper, rename)
         ldr.importer.extension = 'ERR,2'
         ldr.load()
 
+        assert len(deconfigged_helper.viewers['Image'].data_menu.data_labels_loaded) == 2
         assert deconfigged_helper._app._get_assoc_data_parent('Image[ERR,2]') == primary_label
         assert deconfigged_helper._app._get_assoc_data_children(primary_label) == ['Image[ERR,2]']
 
+        # This should raise 'Data labels {unavailable} not able to be loaded...'
+        # if parent child association doesn't work correctly (i.e. the previous behavior)
+        deconfigged_helper.viewers['Image'].data_menu.add_data('Image[ERR,2]')
+
 
 def test_load_image_align_by(deconfigged_helper, image_nddata_wcs):
     ldr = deconfigged_helper.loaders['object']