From 8b997fb0a942ab33f86d20523c719bbbc9d08531 Mon Sep 17 00:00:00 2001
From: Kyle Conroy <kyleconroy@gmail.com>
Date: Fri, 22 Dec 2023 13:15:12 -0500
Subject: [PATCH] helper and viewer-level access to data_labels properties
 (#2626)

* helper and viewer-level access to data_labels properties

* do not expose for table viewers

* test coverage
---
 CHANGES.rst                               |  3 ++
 jdaviz/configs/default/plugins/viewers.py | 53 ++++++++++++++++++-----
 jdaviz/core/helpers.py                    | 12 +++++
 jdaviz/tests/test_user_api.py             |  9 ++++
 4 files changed, 67 insertions(+), 10 deletions(-)

diff --git a/CHANGES.rst b/CHANGES.rst
index 558df7d21a..e81b162f67 100644
--- a/CHANGES.rst
+++ b/CHANGES.rst
@@ -31,6 +31,9 @@ API Changes
 - ``width`` argument in Line Analysis plugin is renamed to ``continuum_width`` and ``width``
   will be removed in a future release. [#2587]
 
+- New API access to ``viz.data_labels``, ``viewer.data_labels_visible``, and
+  ``viewer.data_labels_loaded``. [#2626]
+
 Cubeviz
 ^^^^^^^
 
diff --git a/jdaviz/configs/default/plugins/viewers.py b/jdaviz/configs/default/plugins/viewers.py
index 866378702c..589d589c52 100644
--- a/jdaviz/configs/default/plugins/viewers.py
+++ b/jdaviz/configs/default/plugins/viewers.py
@@ -36,25 +36,58 @@ def __init__(self, *args, **kwargs):
     @property
     def user_api(self):
         # default exposed user APIs.  Can override this method in any particular viewer.
+        if not isinstance(self, TableViewer):
+            expose = ['data_labels', 'data_labels_loaded', 'data_labels_visible']
+        else:
+            expose = []
         if isinstance(self, BqplotImageView):
             if isinstance(self, AstrowidgetsImageViewerMixin):
-                expose = ['save',
-                          'center_on', 'offset_by', 'zoom_level', 'zoom',
-                          'colormap_options', 'set_colormap',
-                          'stretch_options', 'stretch',
-                          'autocut_options', 'cuts',
-                          'marker', 'add_markers', 'remove_markers', 'reset_markers',
-                          'blink_once', 'reset_limits']
+                expose += ['save',
+                           'center_on', 'offset_by', 'zoom_level', 'zoom',
+                           'colormap_options', 'set_colormap',
+                           'stretch_options', 'stretch',
+                           'autocut_options', 'cuts',
+                           'marker', 'add_markers', 'remove_markers', 'reset_markers',
+                           'blink_once', 'reset_limits']
             else:
                 # cubeviz image viewers don't inherit from AstrowidgetsImageViewerMixin yet,
                 # but also shouldn't expose set_limits because of equal aspect ratio concerns
-                expose = []
+                expose += []
         elif isinstance(self, TableViewer):
-            expose = []
+            expose += []
         else:
-            expose = ['set_limits', 'reset_limits']
+            expose += ['set_limits', 'reset_limits']
         return ViewerUserApi(self, expose=expose)
 
+    @property
+    def data_labels_loaded(self):
+        """
+        List of data labels loaded in this viewer.
+
+        Returns
+        -------
+        data_labels : list
+            list of strings
+        """
+        viewer_item = self.jdaviz_app._get_viewer_item(self.reference_id)
+        return [self.jdaviz_app._get_data_item_by_id(data_id)['name']
+                for data_id in viewer_item.get('selected_data_items', {}).keys()]
+
+    @property
+    def data_labels_visible(self):
+        """
+        List of data labels visible in this viewer.
+
+        Returns
+        -------
+        data_labels : list
+            list of strings
+        """
+        viewer_item = self.jdaviz_app._get_viewer_item(self.reference_id)
+        return [self.jdaviz_app._get_data_item_by_id(data_id)['name']
+                for data_id, visibility in viewer_item.get('selected_data_items', {}).items()
+                if visibility != 'hidden']
+
     def reset_limits(self):
         """
         Reset viewer axes limits.
diff --git a/jdaviz/core/helpers.py b/jdaviz/core/helpers.py
index 9ccdc54ade..5948a8c184 100644
--- a/jdaviz/core/helpers.py
+++ b/jdaviz/core/helpers.py
@@ -109,6 +109,18 @@ def load_data(self, data, data_label=None, parser_reference=None, **kwargs):
             kwargs['data_label'] = data_label
         self.app.load_data(data, parser_reference=parser_reference, **kwargs)
 
+    @property
+    def data_labels(self):
+        """
+        List of data labels loaded and available in jdaviz
+
+        Returns
+        -------
+        data_labels : list
+            list of strings
+        """
+        return [data.label for data in self.app.data_collection]
+
     @property
     def plugins(self):
         """
diff --git a/jdaviz/tests/test_user_api.py b/jdaviz/tests/test_user_api.py
index 942b9b29b0..b6bf8c70b5 100644
--- a/jdaviz/tests/test_user_api.py
+++ b/jdaviz/tests/test_user_api.py
@@ -28,3 +28,12 @@ def test_specviz_zoom_level(specviz_helper):
     assert v._obj.state.x_max == 2
     assert v._obj.state.y_min == 1
     assert v._obj.state.y_max == 2
+
+
+def test_specviz_data_labels(specviz_helper, spectrum1d):
+    label = "Test 1D Spectrum"
+    specviz_helper.load_data(spectrum1d, data_label=label)
+
+    assert specviz_helper.data_labels == [label]
+    assert specviz_helper.viewers['spectrum-viewer'].data_labels_loaded == [label]
+    assert specviz_helper.viewers['spectrum-viewer'].data_labels_visible == [label]