From c946d8f4249c54a3a59595f5f6cea339ca79dddd Mon Sep 17 00:00:00 2001
From: Jesse Averbukh <javerbukh@gmail.com>
Date: Thu, 22 Sep 2022 14:01:28 -0400
Subject: [PATCH 1/6] Centralize data label generation

---
 CHANGES.rst                                   |   3 +
 jdaviz/app.py                                 | 250 +++++++++---------
 jdaviz/configs/cubeviz/plugins/parsers.py     |  15 +-
 .../cubeviz/plugins/tests/test_parsers.py     |   8 +-
 .../plugins/model_fitting/model_fitting.py    |   4 +-
 jdaviz/configs/imviz/helper.py                |   4 +-
 jdaviz/configs/imviz/plugins/parsers.py       |  18 +-
 jdaviz/configs/imviz/tests/test_parser.py     |  15 +-
 .../line_analysis/tests/test_lineflux.py      |   5 +-
 jdaviz/configs/specviz/plugins/parsers.py     |  12 +-
 jdaviz/configs/specviz2d/plugins/parsers.py   |   6 +-
 jdaviz/tests/test_app.py                      | 228 ++++++++++++----
 12 files changed, 340 insertions(+), 228 deletions(-)

diff --git a/CHANGES.rst b/CHANGES.rst
index 0125faee1c..3293d8085d 100644
--- a/CHANGES.rst
+++ b/CHANGES.rst
@@ -6,6 +6,9 @@ New Features
 
 - Add support for nonstandard viewer reference names [#1681]
 
+- Centralize data label generation if user does not provide a label with data load. Also
+  prevent duplicate data labels from being added to data collection. [#1672]
+
 Cubeviz
 ^^^^^^^
 
diff --git a/jdaviz/app.py b/jdaviz/app.py
index 31c4210b19..0b19e1d204 100644
--- a/jdaviz/app.py
+++ b/jdaviz/app.py
@@ -2,12 +2,15 @@
 import pathlib
 import re
 import uuid
+import warnings
 from inspect import isclass
 
 from ipywidgets import widget_serialization
 import ipyvue
 
-from astropy.nddata import CCDData
+from astropy.nddata import CCDData, NDData
+from astropy.io import fits
+
 from echo import CallbackProperty, DictCallbackProperty, ListCallbackProperty
 from ipygoldenlayout import GoldenLayout
 from ipysplitpanes import SplitPanes
@@ -403,16 +406,17 @@ def _link_new_data(self, reference_data=None, data_to_be_linked=None):
         ref_data = dc[reference_data] if reference_data else dc[0]
         linked_data = dc[data_to_be_linked] if data_to_be_linked else dc[-1]
 
-        if 'Trace' in linked_data.meta:
-            links = [LinkSame(linked_data.components[1], ref_data.components[0]),
-                     LinkSame(linked_data.components[0], ref_data.components[1])]
-            dc.add_link(links)
-            return
-        elif linked_data.meta.get('Plugin', None) == 'SpectralExtraction':
-            links = [LinkSame(linked_data.components[0], ref_data.components[0]),
-                     LinkSame(linked_data.components[1], ref_data.components[1])]
-            dc.add_link(links)
-            return
+        if linked_data.meta.get('Plugin', None) == 'SpectralExtraction':
+            if 'Trace' in linked_data.meta:
+                links = [LinkSame(linked_data.components[2], ref_data.components[0]),
+                         LinkSame(linked_data.components[0], ref_data.components[1])]
+                dc.add_link(links)
+                return
+            else:
+                links = [LinkSame(linked_data.components[0], ref_data.components[0]),
+                         LinkSame(linked_data.components[1], ref_data.components[1])]
+                dc.add_link(links)
+                return
 
         # The glue-astronomy SpectralCoordinates currently seems incompatible with glue
         # WCSLink. This gets around it until there's an upstream fix.
@@ -490,6 +494,7 @@ def load_data(self, file_obj, parser_reference=None, **kwargs):
         **kwargs : dict
             Additional keywords to be passed into the parser defined by
             ``parser_reference``.
+
         """
         self.loading = True
         try:
@@ -572,7 +577,7 @@ def get_viewer_by_id(self, vid):
         """Like :meth:`get_viewer` but use ID instead of reference name.
         This is useful when reference name is `None`.
         """
-        return self._viewer_store.get(vid)
+        return self._viewer_store[vid]
 
     def get_data_from_viewer(self, viewer_reference, data_label=None,
                              cls='default', include_subsets=True):
@@ -673,8 +678,11 @@ def get_data_from_viewer(self, viewer_reference, data_label=None,
 
                         data[label] = layer_data
 
-        # If a data label was provided, return only the corresponding data, otherwise return all:
-        return data.get(data_label, data)
+        # If a data label was provided, return only the data requested
+        if data_label is not None:
+            return data.get(data_label)
+
+        return data
 
     def get_subsets_from_viewer(self, viewer_reference, data_label=None, subset_type=None):
         """
@@ -896,7 +904,7 @@ def add_data(self, data, data_label=None, notify_done=True):
 
         if not data_label and hasattr(data, "label"):
             data_label = data.label
-        data_label = data_label or "New Data"
+        data_label = self.return_unique_name(data_label)
         self.data_collection[data_label] = data
 
         # Send out a toast message
@@ -905,36 +913,99 @@ def add_data(self, data, data_label=None, notify_done=True):
                 f"Data '{data_label}' successfully added.", sender=self, color="success")
             self.hub.broadcast(snackbar_message)
 
-    @staticmethod
-    def _build_data_label(path, ext=None):
-        """ Build a data label from a filename and data extension
-
-        Builds a data_label out of a filename and data extension.
-        If the input string path already ends with a data
-        extension, the label is returned directly.  Otherwise a valid
-        extension must be specified to append to the file stem.
+    def return_data_label(self, loaded_object, ext=None, alt_name=None, check_unique=True):
+        """
+        Returns a unique data label that can be safely used to load data into data collection.
 
         Parameters
         ----------
-            path : str
-                The data filename to use as a
-            ext : str
-                The data extension to access from the file.
+        loaded_object : str or object
+            The path to a data file or FITS HDUList or image object or Spectrum1D or
+            NDData array or numpy.ndarray.
+        ext : str, optional
+            The extension (or other distinguishing feature) of data used to identify it.
+            For example, "filename[FLUX]" where "FLUX" is the ext value.
+        alt_name : str, optional
+            Alternate names that can be used if none of the options provided are valid.
+        check_unique : bool
+            Included so that this method can be used with data label retrieval in addition
+            to generation.
+
         Returns
         -------
-        A string data label used by Glue
-
-        """
-
-        # if path is not a file or already ends in [ ] extension, assume it is a data-label
-        if not os.path.isfile(path) or re.search(r'(.+)(\[(.*?)\])$', path):
-            return path
-        else:
-            assert ext, 'A data extension must be specified'
-            p = pathlib.Path(path)
-            stem = p.stem.split(os.extsep)[0]
-            label = f'{stem}[{ext}]'
-            return label
+        data_label : str
+            A unique data label that at its root is either given by the user at load time
+            or created by jdaviz using a description of the loaded file's type.
+        """
+        data_label = None
+
+        if loaded_object is None:
+            pass
+        elif isinstance(loaded_object, str):
+            # This is either the full file path or the basename or the user input data label
+            if loaded_object.lower().endswith(('.jpg', '.jpeg', '.png')):
+                file_format = loaded_object.lower().split(".")[-1]
+                loaded_object = os.path.splitext(os.path.basename(loaded_object))[0]
+                data_label = f"{loaded_object}[{file_format}]"
+            elif os.path.isfile(loaded_object) or os.path.isdir(loaded_object):
+                # File that is not an image file
+                data_label = os.path.splitext(os.path.basename(loaded_object))[0]
+            else:
+                # Not a file path so assumed to be a data label
+                data_label = loaded_object
+        elif isinstance(loaded_object, fits.hdu.hdulist.HDUList):
+            if hasattr(loaded_object, 'file_name'):
+                data_label = f"{loaded_object.file_name}[HDU object]"
+            else:
+                data_label = "Unknown HDU object"
+        elif isinstance(loaded_object, Spectrum1D):
+            data_label = "Spectrum1D"
+        elif isinstance(loaded_object, NDData):
+            data_label = "NDData"
+        elif isinstance(loaded_object, np.ndarray):
+            data_label = "Array"
+
+        if data_label is None and alt_name is not None and isinstance(alt_name, str):
+            data_label = alt_name
+        elif data_label is None:
+            data_label = "Unknown"
+
+        if check_unique:
+            data_label = self.return_unique_name(data_label, ext=ext)
+
+        return data_label
+
+    def return_unique_name(self, data_label, ext=None):
+        if data_label is None:
+            data_label = "Unknown"
+
+        # This regex checks for any length of characters that end
+        # with a space followed by parenthesis with a number inside.
+        # If there is a match, the space and parenthesis section will be
+        # removed so that the remainder of the label can be checked
+        # against the data_label.
+        check_if_dup = re.compile(r"(.*)(\s\(\d*\))$")
+        labels = self.data_collection.labels
+        number_of_duplicates = 0
+        for label in labels:
+            # If label is a duplicate of another label
+            if re.fullmatch(check_if_dup, label):
+                label_split = label.split(" ")
+                label_without_dup = " ".join(label_split[:-1])
+                label = label_without_dup
+
+            if ext and f"{data_label}[{ext}]" == label:
+                number_of_duplicates += 1
+            elif ext is None and data_label == label:
+                number_of_duplicates += 1
+
+        if ext:
+            data_label = f"{data_label}[{ext}]"
+
+        if number_of_duplicates > 0:
+            data_label = f"{data_label} ({number_of_duplicates})"
+
+        return data_label
 
     def add_data_to_viewer(self, viewer_reference, data_path,
                            clear_other_data=False, ext=None):
@@ -960,17 +1031,17 @@ def add_data_to_viewer(self, viewer_reference, data_path,
             viewer_item = self._viewer_item_by_id(viewer_reference)
         if viewer_item is None:
             raise ValueError(f"Could not identify viewer with reference {viewer_reference}")
-        data_label = self._build_data_label(data_path, ext=ext)
+        data_label = self.return_data_label(data_path, ext=ext, check_unique=False)
         data_id = self._data_id_from_label(data_label)
 
         if clear_other_data:
-            self._update_selected_data_items(viewer_item.get('id'), {})
+            self._update_selected_data_items(viewer_item['id'], {})
 
-        selected_data_items = viewer_item.get('selected_data_items', {})
+        selected_data_items = viewer_item['selected_data_items']
 
         if data_id is not None:
             selected_data_items[data_id] = 'visible'
-            self._update_selected_data_items(viewer_item.get('id'), selected_data_items)
+            self._update_selected_data_items(viewer_item['id'], selected_data_items)
         else:
             raise ValueError(
                 f"No data item found with label '{data_label}'. Label must be one "
@@ -1006,7 +1077,7 @@ def remove_data_from_viewer(self, viewer_reference, data_path, ext=None):
             is required.
         """
         viewer_item = self._viewer_item_by_reference(viewer_reference)
-        data_label = self._build_data_label(data_path, ext=ext)
+        data_label = self.return_data_label(data_path, ext=ext, check_unique=False)
         data_id = self._data_id_from_label(data_label)
 
         selected_items = viewer_item['selected_data_items']
@@ -1064,55 +1135,7 @@ def get_viewer_ids(self, prefix=None):
     def get_viewer_reference_names(self):
         """Return a list of available viewer reference names."""
         # Cannot sort because of None
-        return [self._viewer_item_by_id(vid).get('reference') for vid in self._viewer_store]
-
-    def _get_first_viewer_reference_name(
-            self, require_no_selected_data=False,
-            require_spectrum_viewer=False,
-            require_spectrum_2d_viewer=False,
-            require_table_viewer=False,
-            require_flux_viewer=False,
-            require_image_viewer=False
-    ):
-        """
-        Return the viewer reference name of the first available viewer.
-        Optionally use ``require_no_selected_data`` to require that the
-        viewer has not yet loaded data, or e.g. ``require_spectrum_viewer``
-        to require that the viewer supports spectrum visualization.
-        """
-        from jdaviz.configs.specviz.plugins.viewers import SpecvizProfileView
-        from jdaviz.configs.specviz2d.plugins import SpectralExtraction
-        from jdaviz.configs.cubeviz.plugins.viewers import CubevizProfileView, CubevizImageView
-        from jdaviz.configs.mosviz.plugins.viewers import (
-            MosvizProfileView, MosvizTableViewer, MosvizProfile2DView
-        )
-
-        spectral_viewers = (SpecvizProfileView, CubevizProfileView, MosvizProfileView)
-        table_viewers = (MosvizTableViewer, )
-        image_viewers = (MosvizProfile2DView, CubevizImageView, SpectralExtraction)
-        flux_viewers = (CubevizImageView, )
-
-        for vid in self._viewer_store:
-            viewer_item = self._viewer_item_by_id(vid)
-            is_returnable = (
-                (require_no_selected_data and not len(viewer_item['selected_data_items'])) or
-                (not require_no_selected_data)
-            )
-            if require_spectrum_viewer:
-                if isinstance(self._viewer_store[vid], spectral_viewers) and is_returnable:
-                    return viewer_item['reference']
-            elif require_table_viewer:
-                if isinstance(self._viewer_store[vid], table_viewers) and is_returnable:
-                    return viewer_item['reference']
-            elif require_image_viewer:
-                if isinstance(self._viewer_store[vid], image_viewers) and is_returnable:
-                    return viewer_item['reference']
-            elif require_flux_viewer:
-                if isinstance(self._viewer_store[vid], flux_viewers) and is_returnable:
-                    return viewer_item['reference']
-            else:
-                if is_returnable:
-                    return viewer_item['reference']
+        return [self._viewer_item_by_id(vid)['reference'] for vid in self._viewer_store]
 
     def _viewer_by_id(self, vid):
         """
@@ -1148,7 +1171,7 @@ def _viewer_item_by_id(self, vid):
         def find_viewer_item(stack_items):
             for stack_item in stack_items:
                 for viewer_item in stack_item.get('viewers'):
-                    if viewer_item.get('id') == vid:
+                    if viewer_item['id'] == vid:
                         return viewer_item
 
                 if len(stack_item.get('children')) > 0:
@@ -1342,9 +1365,6 @@ def _update_selected_data_items(self, viewer_id, selected_items):
         viewer_item = self._viewer_item_by_id(viewer_id)
         viewer = self._viewer_by_id(viewer_id)
 
-        if viewer is None:
-            return
-
         if isinstance(selected_items, list):
             selected_items = {si: 'visible' for si in selected_items}
         # Update the stored selected data items
@@ -1356,8 +1376,14 @@ def _update_selected_data_items(self, viewer_id, selected_items):
 
         # Include any selected data in the viewer
         for data_id, visibility in selected_items.items():
-            label = self._get_data_item_by_id(data_id)['name']
+            data_item = self._get_data_item_by_id(data_id)
 
+            if data_item is None or data_item['name'] is None:
+                warnings.warn(f"No data item with id {data_id} found in "
+                              f"viewer {viewer_id}.")
+                continue
+            else:
+                label = self._get_data_item_by_id(data_id)['name']
             active_data_labels.append(label)
 
             [data] = [x for x in self.data_collection if x.label == label]
@@ -1388,7 +1414,7 @@ def _update_selected_data_items(self, viewer_id, selected_items):
             if data.label not in active_data_labels:
                 viewer.remove_data(data)
                 viewer._layers_with_defaults_applied= [layer_info for layer_info in viewer._layers_with_defaults_applied  # noqa
-                                                       if layer_info['data_label'] != data.label]  # noqa
+                                                       if layer_info['data_label'] != data.label]
                 remove_data_message = RemoveDataMessage(data, viewer,
                                                         viewer_id=viewer_id,
                                                         sender=self)
@@ -1714,29 +1740,7 @@ def compose_viewer_area(viewer_area_items):
 
         for name in config.get('tray', []):
             tray = tray_registry.members.get(name)
-            tray_registry_options = tray.get('viewer_reference_name_kwargs', {})
-
-            # Optional keyword arguments are required to initialize some
-            # tray items. These kwargs specify the viewer reference names that are
-            # assumed to be present in the configuration.
-            optional_tray_kwargs = dict()
-
-            # If viewer reference names need to be passed to the tray item
-            # constructor, pass the names into the constructor in the format
-            # that the tray items expect.
-            for opt_attr, [opt_kwarg, get_name_kwargs] in tray_registry_options.items():
-                opt_value = getattr(
-                    self, opt_attr, self._get_first_viewer_reference_name(**get_name_kwargs)
-                )
-
-                if opt_value is None:
-                    continue
-
-                optional_tray_kwargs[opt_kwarg] = opt_value
-
-            tray_item_instance = tray.get('cls')(
-                app=self, **optional_tray_kwargs
-            )
+            tray_item_instance = tray.get('cls')(app=self)
             tray_item_label = tray.get('label')
 
             self.state.tray_items.append({
diff --git a/jdaviz/configs/cubeviz/plugins/parsers.py b/jdaviz/configs/cubeviz/plugins/parsers.py
index c16d538d27..2ce6d5093b 100644
--- a/jdaviz/configs/cubeviz/plugins/parsers.py
+++ b/jdaviz/configs/cubeviz/plugins/parsers.py
@@ -74,7 +74,7 @@ def parse_data(app, file_obj, data_type=None, data_label=None):
                 for ext, viewer_name in (('SCI', flux_viewer_reference_name),
                                          ('ERR', uncert_viewer_reference_name),
                                          ('DQ', None)):
-                    data_label = f'{file_name}[{ext}]'
+                    data_label = app.return_data_label(file_name, ext)
                     _parse_jwst_s3d(
                         app, hdulist, data_label, ext=ext, viewer_name=viewer_name,
                         flux_viewer_reference_name=flux_viewer_reference_name,
@@ -84,13 +84,12 @@ def parse_data(app, file_obj, data_type=None, data_label=None):
                 for ext, viewer_name in (('DATA', flux_viewer_reference_name),
                                          ('ERR', uncert_viewer_reference_name),
                                          ('QUALITY', None)):
-                    data_label = f'{file_name}[{ext}]'
+                    data_label = app.return_data_label(file_name, ext)
                     _parse_esa_s3d(
                         app, hdulist, data_label, ext=ext, viewer_name=viewer_name,
                         flux_viewer_reference_name=flux_viewer_reference_name,
                         spectrum_viewer_reference_name=spectrum_viewer_reference_name
                     )
-
             else:
                 _parse_hdulist(
                     app, hdulist, file_name=data_label or file_name,
@@ -183,7 +182,7 @@ def _parse_hdulist(app, hdulist, file_name=None,
             continue
 
         is_loaded.append(data_type)
-        data_label = f"{file_name}[{hdu.name}]"
+        data_label = app.return_data_label(file_name, hdu.name)
 
         if data_type == 'flux':
             wcs = WCS(hdu.header, hdulist)
@@ -346,7 +345,7 @@ def _parse_spectrum1d_3d(app, file_obj, data_label=None,
             s1d = Spectrum1D(flux=flux, wcs=file_obj.wcs,
                              meta=standardize_metadata(file_obj.meta))
 
-        cur_data_label = f"{data_label}[{attr.upper()}]"
+        cur_data_label = app.return_data_label(data_label, attr.upper())
         app.add_data(s1d, cur_data_label)
 
         if attr == 'flux':
@@ -361,13 +360,13 @@ def _parse_spectrum1d_3d(app, file_obj, data_label=None,
 
 def _parse_spectrum1d(app, file_obj, data_label=None, spectrum_viewer_reference_name=None):
     if data_label is None:
-        data_label = "Unknown spectrum object"
+        data_label = app.return_data_label(file_obj)
 
     # TODO: glue-astronomy translators only look at the flux property of
     #  specutils Spectrum1D objects. Fix to support uncertainties and masks.
 
-    app.add_data(file_obj, f"{data_label}[FLUX]")
-    app.add_data_to_viewer(spectrum_viewer_reference_name, f"{data_label}[FLUX]")
+    app.add_data(file_obj, data_label)
+    app.add_data_to_viewer(spectrum_viewer_reference_name, data_label)
 
 
 def _get_data_type_by_hdu(hdu):
diff --git a/jdaviz/configs/cubeviz/plugins/tests/test_parsers.py b/jdaviz/configs/cubeviz/plugins/tests/test_parsers.py
index 24b5720729..8f9635f4c2 100644
--- a/jdaviz/configs/cubeviz/plugins/tests/test_parsers.py
+++ b/jdaviz/configs/cubeviz/plugins/tests/test_parsers.py
@@ -14,7 +14,7 @@ def test_fits_image_hdu_parse(image_cube_hdu_obj, cubeviz_helper):
     cubeviz_helper.load_data(image_cube_hdu_obj)
 
     assert len(cubeviz_helper.app.data_collection) == 3
-    assert cubeviz_helper.app.data_collection[0].label.endswith('[FLUX]')
+    assert cubeviz_helper.app.data_collection[0].label == "Unknown HDU object[FLUX]"
 
     # first load should be successful; subsequent attempts should fail
     with pytest.raises(RuntimeError, match=r".?only one (data)?cube.?"):
@@ -88,7 +88,7 @@ def test_fits_image_hdu_parse_from_file(tmpdir, image_cube_hdu_obj, cubeviz_help
     cubeviz_helper.load_data(path)
 
     assert len(cubeviz_helper.app.data_collection) == 3
-    assert cubeviz_helper.app.data_collection[0].label.endswith('[FLUX]')
+    assert cubeviz_helper.app.data_collection[0].label == "test_fits_image.fits[FLUX]"
 
     # This tests the same data as test_fits_image_hdu_parse above.
 
@@ -121,7 +121,7 @@ def test_spectrum3d_parse(image_cube_hdu_obj, cubeviz_helper):
     cubeviz_helper.load_data(sc)
 
     assert len(cubeviz_helper.app.data_collection) == 1
-    assert cubeviz_helper.app.data_collection[0].label.endswith('[FLUX]')
+    assert cubeviz_helper.app.data_collection[0].label == "Unknown spectrum object[FLUX]"
 
     # Same as flux viewer data in test_fits_image_hdu_parse_from_file
     flux_viewer = cubeviz_helper.app.get_viewer(cubeviz_helper._default_flux_viewer_reference_name)
@@ -142,7 +142,7 @@ def test_spectrum1d_parse(spectrum1d, cubeviz_helper):
     cubeviz_helper.load_data(spectrum1d)
 
     assert len(cubeviz_helper.app.data_collection) == 1
-    assert cubeviz_helper.app.data_collection[0].label.endswith('[FLUX]')
+    assert cubeviz_helper.app.data_collection[0].label.endswith('Spectrum1D')
     assert cubeviz_helper.app.data_collection[0].meta['uncertainty_type'] == 'std'
 
     # Coordinate display is only for spatial image, which is missing here.
diff --git a/jdaviz/configs/default/plugins/model_fitting/model_fitting.py b/jdaviz/configs/default/plugins/model_fitting/model_fitting.py
index c18b3ea414..2b46e69b9b 100644
--- a/jdaviz/configs/default/plugins/model_fitting/model_fitting.py
+++ b/jdaviz/configs/default/plugins/model_fitting/model_fitting.py
@@ -760,8 +760,10 @@ def _fit_model_to_cube(self, add_data):
         count = max(map(lambda s: int(next(iter(re.findall(r"\d$", s)), 0)),
                         self.data_collection.labels)) + 1
 
+        label = self.app.return_data_label(f"{self.results_label}[Cube]", ext=count)
+
         # Create new glue data object
-        output_cube = Data(label=f"{self.results_label} [Cube] {count}",
+        output_cube = Data(label=label,
                            coords=data.coords)
         output_cube['flux'] = fitted_spectrum.flux.value
         output_cube.get_component('flux').units = fitted_spectrum.flux.unit.to_string()
diff --git a/jdaviz/configs/imviz/helper.py b/jdaviz/configs/imviz/helper.py
index 10bdfd0a04..916dbd5eb5 100644
--- a/jdaviz/configs/imviz/helper.py
+++ b/jdaviz/configs/imviz/helper.py
@@ -142,7 +142,7 @@ def load_data(self, data, data_label=None, do_link=True, show_in_viewer=True, **
 
                 # This will only overwrite if not provided.
                 if not data_label:
-                    kw['data_label'] = cur_data_label
+                    kw['data_label'] = None
                 else:
                     kw['data_label'] = data_label
 
@@ -165,7 +165,7 @@ def load_data(self, data, data_label=None, do_link=True, show_in_viewer=True, **
 
                 # This will only append index to data label if provided.
                 if data_label:
-                    kw['data_label'] = f'{data_label}_{i}'
+                    kw['data_label'] = data_label
 
                 self.app.load_data(data[i, :, :], parser_reference='imviz-data-parser', **kw)
 
diff --git a/jdaviz/configs/imviz/plugins/parsers.py b/jdaviz/configs/imviz/plugins/parsers.py
index dd57582431..ba7a8d6513 100644
--- a/jdaviz/configs/imviz/plugins/parsers.py
+++ b/jdaviz/configs/imviz/plugins/parsers.py
@@ -1,6 +1,4 @@
-import base64
 import os
-import uuid
 
 import numpy as np
 from asdf.fits_embed import AsdfInFits
@@ -56,8 +54,6 @@ def parse_data(app, file_obj, ext=None, data_label=None):
             with fits.open(file_obj) as pf:
                 _parse_image(app, pf, data_label, ext=ext)
     else:
-        if data_label is None:
-            data_label = f'imviz_data|{str(base64.b85encode(uuid.uuid4().bytes), "utf-8")}'
         _parse_image(app, file_obj, data_label, ext=ext)
 
 
@@ -121,20 +117,14 @@ def get_image_data_iterator(app, file_obj, data_label, ext=None):
 
 
 def _parse_image(app, file_obj, data_label, ext=None):
-    from jdaviz.core.helpers import _next_subset_num
-
+    if app is None:
+        raise ValueError("app is None, cannot proceed")
     if data_label is None:
-        raise NotImplementedError('data_label should be set by now')
-
+        data_label = app.return_data_label(file_obj, ext, alt_name="image_data")
     data_iter = get_image_data_iterator(app, file_obj, data_label, ext=ext)
 
     for data, data_label in data_iter:
-
-        # avoid duplicate data labels in collection
-        if data_label in app.data_collection.labels:
-            i = _next_subset_num(data_label, app.data_collection)
-            data_label = f'{data_label} {i}'
-
+        data_label = app.return_data_label(data_label, alt_name="image_data")
         app.add_data(data, data_label)
 
     # Do not run link_image_data here. We do it at the end in Imviz.load_data()
diff --git a/jdaviz/configs/imviz/tests/test_parser.py b/jdaviz/configs/imviz/tests/test_parser.py
index 074dd698b6..c348d63b98 100644
--- a/jdaviz/configs/imviz/tests/test_parser.py
+++ b/jdaviz/configs/imviz/tests/test_parser.py
@@ -63,7 +63,7 @@ def setup_class(self):
         self.jwst_asdf_url_2 = 'https://stsci.box.com/shared/static/d5k9z5j05dgfv6ljgie483w21kmpevni.fits'  # noqa: E501
 
     def test_no_data_label(self):
-        with pytest.raises(NotImplementedError, match='should be set'):
+        with pytest.raises(ValueError, match='app is None'):
             _parse_image(None, None, None, False)
 
     def test_hdulist_no_image(self, imviz_helper):
@@ -94,13 +94,14 @@ def test_parse_numpy_array_3d(self, imviz_helper, manual_loop):
         n_slices = 5
         slice_shape = (2, 3)
         arr = np.stack([np.zeros(slice_shape) + i for i in range(n_slices)])
+        data_label = 'my_slices'
 
         if not manual_loop:
             # We use higher level load_data() here to make sure linking does not crash.
-            imviz_helper.load_data(arr, data_label='my_slices')
+            imviz_helper.load_data(arr, data_label=data_label)
         else:
             for i in range(n_slices):
-                imviz_helper.load_data(arr[i, :, :], data_label=f'my_slices_{i}', do_link=False)
+                imviz_helper.load_data(arr[i, :, :], data_label='my_slices', do_link=False)
             imviz_helper.link_data(error_on_fail=True)
 
         assert len(imviz_helper.app.data_collection) == n_slices
@@ -109,7 +110,7 @@ def test_parse_numpy_array_3d(self, imviz_helper, manual_loop):
         for i in range(n_slices):
             data = imviz_helper.app.data_collection[i]
             comp = data.get_component('DATA')
-            assert data.label == f'my_slices_{i}'
+            assert data.label == (f'my_slices ({i})' if i > 0 else 'my_slices')
             assert data.shape == slice_shape
             assert_array_equal(comp.data, i)
 
@@ -336,7 +337,7 @@ def test_parse_jwst_nircam_level2(self, imviz_helper):
 
             imviz_helper.load_data(pf, ext='SCI', data_label='TEST', show_in_viewer=False)
             data = imviz_helper.app.data_collection[1]
-            assert data.label.endswith('[DATA] 1')
+            assert data.label.endswith('[DATA] (1)')
 
             # Load all extensions
             imviz_helper.app.data_collection.clear()
@@ -453,14 +454,14 @@ def test_parse_hst_drz(self, imviz_helper):
             # Default behavior: Load first image
             imviz_helper.load_data(pf, show_in_viewer=False)
             data = imviz_helper.app.data_collection[3]
-            assert data.label.startswith('imviz_data|') and data.label.endswith('[SCI,1]')
+            assert data.label.startswith('Unknown HDU object') and data.label.endswith('[SCI,1]')
             assert_allclose(data.meta['PHOTFLAM'], 7.8711728E-20)
             assert 'SCI,1' in data.components
 
             # Request specific extension (name only), use given label
             imviz_helper.load_data(pf, ext='CTX', show_in_viewer=False)
             data = imviz_helper.app.data_collection[4]
-            assert data.label.startswith('imviz_data|') and data.label.endswith('[CTX,1]')
+            assert data.label.startswith('Unknown HDU object') and data.label.endswith('[CTX,1]')
             assert data.meta['EXTNAME'] == 'CTX'
             assert 'CTX,1' in data.components
 
diff --git a/jdaviz/configs/specviz/plugins/line_analysis/tests/test_lineflux.py b/jdaviz/configs/specviz/plugins/line_analysis/tests/test_lineflux.py
index 947be1f1c7..7ec76cd0aa 100644
--- a/jdaviz/configs/specviz/plugins/line_analysis/tests/test_lineflux.py
+++ b/jdaviz/configs/specviz/plugins/line_analysis/tests/test_lineflux.py
@@ -83,9 +83,8 @@ def test_cubeviz_collapse_fluxunits(spectrum1d_cube_custom_fluxunit, spectra_flu
         cubeviz_helper.app.get_viewer('spectrum-viewer').state.function = function
 
         lineflux_result = _calculate_line_flux(cubeviz_helper)
-
-        autocollapsed_spectrum_unit = cubeviz_helper.specviz.get_spectra()[data_label + "[FLUX]"
-                                                                           ].flux.unit
+        autocollapsed_spectrum_unit = (cubeviz_helper.
+                                       specviz.get_spectra()[f"{data_label}[FLUX]"].flux.unit)
         # Futureproofing: Eventually Cubeviz autocollapse will change the flux units of the
         # spectra depending on whether the spectrum was collapsed OVER SPAXELS or not. Only
         # do the assertion check if we KNOW what the expected lineflux results should be
diff --git a/jdaviz/configs/specviz/plugins/parsers.py b/jdaviz/configs/specviz/plugins/parsers.py
index c82c870f5c..11c7ec7fa6 100644
--- a/jdaviz/configs/specviz/plugins/parsers.py
+++ b/jdaviz/configs/specviz/plugins/parsers.py
@@ -1,6 +1,4 @@
-import base64
 import pathlib
-import uuid
 
 import numpy as np
 from astropy.io.registry import IORegistryError
@@ -31,18 +29,16 @@ def specviz_spectrum1d_parser(app, data, data_label=None, format=None, show_in_v
         Reference name for the viewer
     """
     spectrum_viewer_reference_name = app._jdaviz_helper._default_spectrum_viewer_reference_name
-    # If no data label is assigned, give it a unique identifier
+    # If no data label is assigned, give it a unique name
     if not data_label:
-        data_label = "specviz_data|" + str(
-            base64.b85encode(uuid.uuid4().bytes), "utf-8"
-        )
+        data_label = app.return_data_label(data, alt_name="specviz_data")
 
     if isinstance(data, SpectrumCollection):
         raise TypeError("SpectrumCollection detected."
                         " Please provide a Spectrum1D or SpectrumList")
     elif isinstance(data, Spectrum1D):
         data = [data]
-        data_label = [data_label]
+        data_label = [app.return_data_label(data_label, alt_name="specviz_data")]
     # No special processing is needed in this case, but we include it for completeness
     elif isinstance(data, SpectrumList):
         pass
@@ -54,7 +50,7 @@ def specviz_spectrum1d_parser(app, data, data_label=None, format=None, show_in_v
         if path.is_file():
             try:
                 data = [Spectrum1D.read(str(path), format=format)]
-                data_label = [data_label]
+                data_label = [app.return_data_label(data_label, alt_name="specviz_data")]
             except IORegistryError:
                 # Multi-extension files may throw a registry error
                 data = SpectrumList.read(str(path), format=format)
diff --git a/jdaviz/configs/specviz2d/plugins/parsers.py b/jdaviz/configs/specviz2d/plugins/parsers.py
index 6fc0a8a155..199e0fd622 100644
--- a/jdaviz/configs/specviz2d/plugins/parsers.py
+++ b/jdaviz/configs/specviz2d/plugins/parsers.py
@@ -64,7 +64,11 @@ def spec2d_1d_parser(app, data_obj, data_label=None, show_in_viewer=True):
 
         data_obj = Spectrum1D(flux, spectral_axis=spectral_axis, meta=metadata)
 
-    app.data_collection[data_label] = data_obj
+        data_label = app.return_data_label(data_label, alt_name="specviz2d_data")
+        app.data_collection[data_label] = data_obj
+
+    else:
+        raise NotImplementedError("Spectrum2d parser only takes a filename")
 
     if show_in_viewer:
         app.add_data_to_viewer(
diff --git a/jdaviz/tests/test_app.py b/jdaviz/tests/test_app.py
index 7531f8d5c7..1d6c48e7d0 100644
--- a/jdaviz/tests/test_app.py
+++ b/jdaviz/tests/test_app.py
@@ -1,57 +1,171 @@
-import pytest
-
-from jdaviz import Application, Specviz
-from jdaviz.configs.default.plugins.gaussian_smooth.gaussian_smooth import GaussianSmooth
-
-
-# This applies to all viz but testing with Imviz should be enough.
-def test_viewer_calling_app(imviz_helper):
-    viewer = imviz_helper.default_viewer
-    assert viewer.session.jdaviz_app is imviz_helper.app
-
-
-def test_get_tray_item_from_name():
-    app = Application(configuration='default')
-    plg = app.get_tray_item_from_name('g-gaussian-smooth')
-    assert isinstance(plg, GaussianSmooth)
-
-    with pytest.raises(KeyError, match='not found in app'):
-        app.get_tray_item_from_name('imviz-compass')
-
-
-def test_nonstandard_specviz_viewer_name(spectrum1d):
-    config = {'settings': {'configuration': 'nonstandard',
-                           'data': {'parser': 'specviz-spectrum1d-parser'},
-                           'visible': {'menu_bar': False,
-                                       'toolbar': True,
-                                       'tray': True,
-                                       'tab_headers': False},
-                           'context': {'notebook': {'max_height': '750px'}}},
-              'toolbar': ['g-data-tools', 'g-subset-tools'],
-              'tray': ['g-metadata-viewer',
-                       'g-plot-options',
-                       'g-subset-plugin',
-                       'g-gaussian-smooth',
-                       'g-model-fitting',
-                       'g-unit-conversion',
-                       'g-line-list',
-                       'specviz-line-analysis',
-                       'g-export-plot'],
-              'viewer_area': [{'container': 'col',
-                               'children': [{'container': 'row',
-                                             'viewers': [{'name': 'H',
-                                                          'plot': 'specviz-profile-viewer',
-                                                          'reference': 'h'},
-                                                         {'name': 'K',
-                                                          'plot': 'specviz-profile-viewer',
-                                                          'reference': 'k'}]}]}]}
-
-    class Customviz(Specviz):
-        _default_configuration = config
-        _default_spectrum_viewer_reference_name = 'h'
-
-    viz = Customviz()
-    assert viz.app.get_viewer_reference_names() == ['h', 'k']
-
-    viz.load_spectrum(spectrum1d, data_label='example label')
-    assert not len(viz.app.get_data_from_viewer("h", "non-existent label"))
+import pytest
+
+from jdaviz import Application, Specviz
+from jdaviz.configs.default.plugins.gaussian_smooth.gaussian_smooth import GaussianSmooth
+
+
+# This applies to all viz but testing with Imviz should be enough.
+def test_viewer_calling_app(imviz_helper):
+    viewer = imviz_helper.default_viewer
+    assert viewer.session.jdaviz_app is imviz_helper.app
+
+
+def test_get_tray_item_from_name():
+    app = Application(configuration='default')
+    plg = app.get_tray_item_from_name('g-gaussian-smooth')
+    assert isinstance(plg, GaussianSmooth)
+
+    with pytest.raises(KeyError, match='not found in app'):
+        app.get_tray_item_from_name('imviz-compass')
+
+
+def test_nonstandard_specviz_viewer_name(spectrum1d):
+    config = {'settings': {'configuration': 'nonstandard',
+                           'data': {'parser': 'specviz-spectrum1d-parser'},
+                           'visible': {'menu_bar': False,
+                                       'toolbar': True,
+                                       'tray': True,
+                                       'tab_headers': False},
+                           'context': {'notebook': {'max_height': '750px'}}},
+              'toolbar': ['g-data-tools', 'g-subset-tools'],
+              'tray': ['g-metadata-viewer',
+                       'g-plot-options',
+                       'g-subset-plugin',
+                       'g-gaussian-smooth',
+                       'g-model-fitting',
+                       'g-unit-conversion',
+                       'g-line-list',
+                       'specviz-line-analysis',
+                       'g-export-plot'],
+              'viewer_area': [{'container': 'col',
+                               'children': [{'container': 'row',
+                                             'viewers': [{'name': 'H',
+                                                          'plot': 'specviz-profile-viewer',
+                                                          'reference': 'h'},
+                                                         {'name': 'K',
+                                                          'plot': 'specviz-profile-viewer',
+                                                          'reference': 'k'}]}]}]}
+
+    class Customviz(Specviz):
+        _default_configuration = config
+        _default_spectrum_viewer_reference_name = 'h'
+
+    viz = Customviz()
+    assert viz.app.get_viewer_reference_names() == ['h', 'k']
+
+    viz.load_spectrum(spectrum1d, data_label='example label')
+    assert not len(viz.app.get_data_from_viewer("h", "non-existent label"))
+
+
+# This applies to all viz but testing with Imviz should be enough.
+def test_viewer_calling_app(imviz_helper):
+    viewer = imviz_helper.default_viewer
+    assert viewer.session.jdaviz_app is imviz_helper.app
+
+
+def test_get_tray_item_from_name():
+    app = Application(configuration='default')
+    plg = app.get_tray_item_from_name('g-gaussian-smooth')
+    assert isinstance(plg, GaussianSmooth)
+
+    with pytest.raises(KeyError, match='not found in app'):
+        app.get_tray_item_from_name('imviz-compass')
+
+
+# This applies to all viz but testing with Imviz should be enough.
+def test_viewer_calling_app(imviz_helper):
+    viewer = imviz_helper.default_viewer
+    assert viewer.session.jdaviz_app is imviz_helper.app
+
+
+def test_get_tray_item_from_name():
+    app = Application(configuration='default')
+    plg = app.get_tray_item_from_name('g-gaussian-smooth')
+    assert isinstance(plg, GaussianSmooth)
+
+    with pytest.raises(KeyError, match='not found in app'):
+        app.get_tray_item_from_name('imviz-compass')
+
+
+def test_duplicate_data_labels(specviz_helper, spectrum1d):
+    specviz_helper.load_spectrum(spectrum1d, data_label="test")
+    specviz_helper.load_spectrum(spectrum1d, data_label="test")
+    dc = specviz_helper.app.data_collection
+    assert dc[0].label == "test"
+    assert dc[1].label == "test (1)"
+    specviz_helper.load_spectrum(spectrum1d, data_label="test_1")
+    specviz_helper.load_spectrum(spectrum1d, data_label="test")
+    assert dc[2].label == "test_1"
+    assert dc[3].label == "test (2)"
+
+
+def test_duplicate_data_labels_with_brackets(specviz_helper, spectrum1d):
+    specviz_helper.load_spectrum(spectrum1d, data_label="test[test]")
+    specviz_helper.load_spectrum(spectrum1d, data_label="test[test]")
+    dc = specviz_helper.app.data_collection
+    assert len(dc) == 2
+    assert dc[0].label == "test[test]"
+    assert dc[1].label == "test[test] (1)"
+
+
+def test_return_data_label_is_none(specviz_helper):
+    data_label = specviz_helper.app.return_data_label(None)
+    assert data_label == "Unknown"
+
+
+def test_return_data_label_is_image(specviz_helper):
+    data_label = specviz_helper.app.return_data_label("data/path/test.jpg")
+    assert data_label == "test[jpg]"
+
+
+def test_hdulist_with_filename(cubeviz_helper, image_cube_hdu_obj):
+    image_cube_hdu_obj.file_name = "test"
+    data_label = cubeviz_helper.app.return_data_label(image_cube_hdu_obj)
+    assert data_label == "test[HDU object]"
+
+
+@pytest.mark.filterwarnings('ignore:.* is a low contrast image')
+def test_file_path_not_image(imviz_helper, tmp_path):
+    path = tmp_path / "myimage.fits"
+    path.touch()
+    data_label = imviz_helper.app.return_data_label(str(path))
+    assert data_label == "myimage"
+
+
+def test_unique_name_variations(specviz_helper, spectrum1d):
+    data_label = specviz_helper.app.return_unique_name(None)
+    assert data_label == "Unknown"
+
+    specviz_helper.load_spectrum(spectrum1d, data_label="test[flux]")
+    data_label = specviz_helper.app.return_data_label("test[flux]", ext="flux")
+    assert data_label == "test[flux][flux]"
+
+    data_label = specviz_helper.app.return_data_label("test", ext="flux")
+    assert data_label == "test[flux] (1)"
+
+
+def test_substring_in_label(specviz_helper, spectrum1d):
+    specviz_helper.load_spectrum(spectrum1d, data_label="M31")
+    specviz_helper.load_spectrum(spectrum1d, data_label="M32")
+    data_label = specviz_helper.app.return_data_label("M")
+    assert data_label == "M"
+
+
+@pytest.mark.parametrize('data_label', ('111111', 'aaaaa', '///(#$@)',
+                                        'two  spaces  repeating',
+                                        'word42word42word  two  spaces'))
+def test_edge_cases(specviz_helper, spectrum1d, data_label):
+    dc = specviz_helper.app.data_collection
+
+    specviz_helper.load_spectrum(spectrum1d, data_label=data_label)
+    specviz_helper.load_spectrum(spectrum1d, data_label=data_label)
+    assert dc[1].label == f"{data_label} (1)"
+
+    specviz_helper.load_spectrum(spectrum1d, data_label=data_label)
+    assert dc[2].label == f"{data_label} (2)"
+
+
+def test_case_that_breaks_return_label(specviz_helper, spectrum1d):
+    specviz_helper.load_spectrum(spectrum1d, data_label="this will break (1)")
+    with pytest.raises(UserWarning, match='No data item'):
+        specviz_helper.load_spectrum(spectrum1d, data_label="this will break")

From b15c631615713d303c7096fd8ceefb6d7e488643 Mon Sep 17 00:00:00 2001
From: Jesse Averbukh <javerbukh@gmail.com>
Date: Tue, 11 Oct 2022 08:54:54 -0400
Subject: [PATCH 2/6] Remove unused tests

---
 jdaviz/tests/test_app.py | 30 ------------------------------
 1 file changed, 30 deletions(-)

diff --git a/jdaviz/tests/test_app.py b/jdaviz/tests/test_app.py
index 1d6c48e7d0..7dee11a24e 100644
--- a/jdaviz/tests/test_app.py
+++ b/jdaviz/tests/test_app.py
@@ -57,36 +57,6 @@ class Customviz(Specviz):
     assert not len(viz.app.get_data_from_viewer("h", "non-existent label"))
 
 
-# This applies to all viz but testing with Imviz should be enough.
-def test_viewer_calling_app(imviz_helper):
-    viewer = imviz_helper.default_viewer
-    assert viewer.session.jdaviz_app is imviz_helper.app
-
-
-def test_get_tray_item_from_name():
-    app = Application(configuration='default')
-    plg = app.get_tray_item_from_name('g-gaussian-smooth')
-    assert isinstance(plg, GaussianSmooth)
-
-    with pytest.raises(KeyError, match='not found in app'):
-        app.get_tray_item_from_name('imviz-compass')
-
-
-# This applies to all viz but testing with Imviz should be enough.
-def test_viewer_calling_app(imviz_helper):
-    viewer = imviz_helper.default_viewer
-    assert viewer.session.jdaviz_app is imviz_helper.app
-
-
-def test_get_tray_item_from_name():
-    app = Application(configuration='default')
-    plg = app.get_tray_item_from_name('g-gaussian-smooth')
-    assert isinstance(plg, GaussianSmooth)
-
-    with pytest.raises(KeyError, match='not found in app'):
-        app.get_tray_item_from_name('imviz-compass')
-
-
 def test_duplicate_data_labels(specviz_helper, spectrum1d):
     specviz_helper.load_spectrum(spectrum1d, data_label="test")
     specviz_helper.load_spectrum(spectrum1d, data_label="test")

From 2e078d680823186463a719c77b7d43694290ef8a Mon Sep 17 00:00:00 2001
From: Jesse Averbukh <javerbukh@gmail.com>
Date: Tue, 11 Oct 2022 09:07:03 -0400
Subject: [PATCH 3/6] Re-add code that was removed during rebase

---
 jdaviz/app.py | 50 +++++++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 49 insertions(+), 1 deletion(-)

diff --git a/jdaviz/app.py b/jdaviz/app.py
index 0b19e1d204..b9dd70347b 100644
--- a/jdaviz/app.py
+++ b/jdaviz/app.py
@@ -1135,7 +1135,55 @@ def get_viewer_ids(self, prefix=None):
     def get_viewer_reference_names(self):
         """Return a list of available viewer reference names."""
         # Cannot sort because of None
-        return [self._viewer_item_by_id(vid)['reference'] for vid in self._viewer_store]
+        return [self._viewer_item_by_id(vid).get('reference') for vid in self._viewer_store]
+
+    def _get_first_viewer_reference_name(
+            self, require_no_selected_data=False,
+            require_spectrum_viewer=False,
+            require_spectrum_2d_viewer=False,
+            require_table_viewer=False,
+            require_flux_viewer=False,
+            require_image_viewer=False
+    ):
+        """
+        Return the viewer reference name of the first available viewer.
+        Optionally use ``require_no_selected_data`` to require that the
+        viewer has not yet loaded data, or e.g. ``require_spectrum_viewer``
+        to require that the viewer supports spectrum visualization.
+        """
+        from jdaviz.configs.specviz.plugins.viewers import SpecvizProfileView
+        from jdaviz.configs.specviz2d.plugins import SpectralExtraction
+        from jdaviz.configs.cubeviz.plugins.viewers import CubevizProfileView, CubevizImageView
+        from jdaviz.configs.mosviz.plugins.viewers import (
+            MosvizProfileView, MosvizTableViewer, MosvizProfile2DView
+        )
+
+        spectral_viewers = (SpecvizProfileView, CubevizProfileView, MosvizProfileView)
+        table_viewers = (MosvizTableViewer, )
+        image_viewers = (MosvizProfile2DView, CubevizImageView, SpectralExtraction)
+        flux_viewers = (CubevizImageView, )
+
+        for vid in self._viewer_store:
+            viewer_item = self._viewer_item_by_id(vid)
+            is_returnable = (
+                (require_no_selected_data and not len(viewer_item['selected_data_items'])) or
+                (not require_no_selected_data)
+            )
+            if require_spectrum_viewer:
+                if isinstance(self._viewer_store[vid], spectral_viewers) and is_returnable:
+                    return viewer_item['reference']
+            elif require_table_viewer:
+                if isinstance(self._viewer_store[vid], table_viewers) and is_returnable:
+                    return viewer_item['reference']
+            elif require_image_viewer:
+                if isinstance(self._viewer_store[vid], image_viewers) and is_returnable:
+                    return viewer_item['reference']
+            elif require_flux_viewer:
+                if isinstance(self._viewer_store[vid], flux_viewers) and is_returnable:
+                    return viewer_item['reference']
+            else:
+                if is_returnable:
+                    return viewer_item['reference']
 
     def _viewer_by_id(self, vid):
         """

From b0f76dde11792a80749988908121821ea4c3a637 Mon Sep 17 00:00:00 2001
From: Jesse Averbukh <javerbukh@gmail.com>
Date: Tue, 11 Oct 2022 10:26:44 -0400
Subject: [PATCH 4/6] Undo rebase errors

---
 jdaviz/app.py | 65 ++++++++++++++++++++++++++++++++++-----------------
 1 file changed, 43 insertions(+), 22 deletions(-)

diff --git a/jdaviz/app.py b/jdaviz/app.py
index b9dd70347b..0d82c45ad8 100644
--- a/jdaviz/app.py
+++ b/jdaviz/app.py
@@ -406,17 +406,16 @@ def _link_new_data(self, reference_data=None, data_to_be_linked=None):
         ref_data = dc[reference_data] if reference_data else dc[0]
         linked_data = dc[data_to_be_linked] if data_to_be_linked else dc[-1]
 
-        if linked_data.meta.get('Plugin', None) == 'SpectralExtraction':
-            if 'Trace' in linked_data.meta:
-                links = [LinkSame(linked_data.components[2], ref_data.components[0]),
-                         LinkSame(linked_data.components[0], ref_data.components[1])]
-                dc.add_link(links)
-                return
-            else:
-                links = [LinkSame(linked_data.components[0], ref_data.components[0]),
-                         LinkSame(linked_data.components[1], ref_data.components[1])]
-                dc.add_link(links)
-                return
+        if 'Trace' in linked_data.meta:
+            links = [LinkSame(linked_data.components[1], ref_data.components[0]),
+                     LinkSame(linked_data.components[0], ref_data.components[1])]
+            dc.add_link(links)
+            return
+        elif linked_data.meta.get('Plugin', None) == 'SpectralExtraction':
+            links = [LinkSame(linked_data.components[0], ref_data.components[0]),
+                     LinkSame(linked_data.components[1], ref_data.components[1])]
+            dc.add_link(links)
+            return
 
         # The glue-astronomy SpectralCoordinates currently seems incompatible with glue
         # WCSLink. This gets around it until there's an upstream fix.
@@ -577,7 +576,7 @@ def get_viewer_by_id(self, vid):
         """Like :meth:`get_viewer` but use ID instead of reference name.
         This is useful when reference name is `None`.
         """
-        return self._viewer_store[vid]
+        return self._viewer_store.get(vid)
 
     def get_data_from_viewer(self, viewer_reference, data_label=None,
                              cls='default', include_subsets=True):
@@ -678,11 +677,8 @@ def get_data_from_viewer(self, viewer_reference, data_label=None,
 
                         data[label] = layer_data
 
-        # If a data label was provided, return only the data requested
-        if data_label is not None:
-            return data.get(data_label)
-
-        return data
+        # If a data label was provided, return only the corresponding data, otherwise return all:
+        return data.get(data_label, data)
 
     def get_subsets_from_viewer(self, viewer_reference, data_label=None, subset_type=None):
         """
@@ -1037,11 +1033,11 @@ def add_data_to_viewer(self, viewer_reference, data_path,
         if clear_other_data:
             self._update_selected_data_items(viewer_item['id'], {})
 
-        selected_data_items = viewer_item['selected_data_items']
+        selected_data_items = viewer_item.get('selected_data_items', {})
 
         if data_id is not None:
             selected_data_items[data_id] = 'visible'
-            self._update_selected_data_items(viewer_item['id'], selected_data_items)
+            self._update_selected_data_items(viewer_item.get('id'), selected_data_items)
         else:
             raise ValueError(
                 f"No data item found with label '{data_label}'. Label must be one "
@@ -1219,7 +1215,7 @@ def _viewer_item_by_id(self, vid):
         def find_viewer_item(stack_items):
             for stack_item in stack_items:
                 for viewer_item in stack_item.get('viewers'):
-                    if viewer_item['id'] == vid:
+                    if viewer_item.get('id') == vid:
                         return viewer_item
 
                 if len(stack_item.get('children')) > 0:
@@ -1413,6 +1409,9 @@ def _update_selected_data_items(self, viewer_id, selected_items):
         viewer_item = self._viewer_item_by_id(viewer_id)
         viewer = self._viewer_by_id(viewer_id)
 
+        if viewer is None:
+            return
+
         if isinstance(selected_items, list):
             selected_items = {si: 'visible' for si in selected_items}
         # Update the stored selected data items
@@ -1462,7 +1461,7 @@ def _update_selected_data_items(self, viewer_id, selected_items):
             if data.label not in active_data_labels:
                 viewer.remove_data(data)
                 viewer._layers_with_defaults_applied= [layer_info for layer_info in viewer._layers_with_defaults_applied  # noqa
-                                                       if layer_info['data_label'] != data.label]
+                                                       if layer_info['data_label'] != data.label]  # noqa
                 remove_data_message = RemoveDataMessage(data, viewer,
                                                         viewer_id=viewer_id,
                                                         sender=self)
@@ -1788,7 +1787,29 @@ def compose_viewer_area(viewer_area_items):
 
         for name in config.get('tray', []):
             tray = tray_registry.members.get(name)
-            tray_item_instance = tray.get('cls')(app=self)
+            tray_registry_options = tray.get('viewer_reference_name_kwargs', {})
+
+            # Optional keyword arguments are required to initialize some
+            # tray items. These kwargs specify the viewer reference names that are
+            # assumed to be present in the configuration.
+            optional_tray_kwargs = dict()
+
+            # If viewer reference names need to be passed to the tray item
+            # constructor, pass the names into the constructor in the format
+            # that the tray items expect.
+            for opt_attr, [opt_kwarg, get_name_kwargs] in tray_registry_options.items():
+                opt_value = getattr(
+                    self, opt_attr, self._get_first_viewer_reference_name(**get_name_kwargs)
+                )
+
+                if opt_value is None:
+                    continue
+
+                optional_tray_kwargs[opt_kwarg] = opt_value
+
+            tray_item_instance = tray.get('cls')(
+                app=self, **optional_tray_kwargs
+            )
             tray_item_label = tray.get('label')
 
             self.state.tray_items.append({

From b1f0672d68dcb3112105439ba2d17c0686248df7 Mon Sep 17 00:00:00 2001
From: Jesse Averbukh <javerbukh@gmail.com>
Date: Tue, 11 Oct 2022 13:00:34 -0400
Subject: [PATCH 5/6] Address review comments

---
 jdaviz/app.py                             | 4 ++--
 jdaviz/configs/imviz/tests/test_parser.py | 2 +-
 jdaviz/tests/test_app.py                  | 1 -
 3 files changed, 3 insertions(+), 4 deletions(-)

diff --git a/jdaviz/app.py b/jdaviz/app.py
index 0d82c45ad8..e4f52613ea 100644
--- a/jdaviz/app.py
+++ b/jdaviz/app.py
@@ -931,7 +931,7 @@ def return_data_label(self, loaded_object, ext=None, alt_name=None, check_unique
         -------
         data_label : str
             A unique data label that at its root is either given by the user at load time
-            or created by jdaviz using a description of the loaded file's type.
+            or created by Jdaviz using a description of the loaded filetype.
         """
         data_label = None
 
@@ -1031,7 +1031,7 @@ def add_data_to_viewer(self, viewer_reference, data_path,
         data_id = self._data_id_from_label(data_label)
 
         if clear_other_data:
-            self._update_selected_data_items(viewer_item['id'], {})
+            self._update_selected_data_items(viewer_item.get('id'), {})
 
         selected_data_items = viewer_item.get('selected_data_items', {})
 
diff --git a/jdaviz/configs/imviz/tests/test_parser.py b/jdaviz/configs/imviz/tests/test_parser.py
index c348d63b98..8b0f38da88 100644
--- a/jdaviz/configs/imviz/tests/test_parser.py
+++ b/jdaviz/configs/imviz/tests/test_parser.py
@@ -101,7 +101,7 @@ def test_parse_numpy_array_3d(self, imviz_helper, manual_loop):
             imviz_helper.load_data(arr, data_label=data_label)
         else:
             for i in range(n_slices):
-                imviz_helper.load_data(arr[i, :, :], data_label='my_slices', do_link=False)
+                imviz_helper.load_data(arr[i, :, :], data_label=data_label, do_link=False)
             imviz_helper.link_data(error_on_fail=True)
 
         assert len(imviz_helper.app.data_collection) == n_slices
diff --git a/jdaviz/tests/test_app.py b/jdaviz/tests/test_app.py
index 7dee11a24e..611b5a173d 100644
--- a/jdaviz/tests/test_app.py
+++ b/jdaviz/tests/test_app.py
@@ -94,7 +94,6 @@ def test_hdulist_with_filename(cubeviz_helper, image_cube_hdu_obj):
     assert data_label == "test[HDU object]"
 
 
-@pytest.mark.filterwarnings('ignore:.* is a low contrast image')
 def test_file_path_not_image(imviz_helper, tmp_path):
     path = tmp_path / "myimage.fits"
     path.touch()

From e6b515a50eb9d4711a39439eff2a3339e6ea3122 Mon Sep 17 00:00:00 2001
From: Jesse Averbukh <javerbukh@gmail.com>
Date: Tue, 11 Oct 2022 13:02:09 -0400
Subject: [PATCH 6/6] Remove comment

---
 jdaviz/configs/imviz/helper.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/jdaviz/configs/imviz/helper.py b/jdaviz/configs/imviz/helper.py
index 916dbd5eb5..f70c3f4981 100644
--- a/jdaviz/configs/imviz/helper.py
+++ b/jdaviz/configs/imviz/helper.py
@@ -163,7 +163,6 @@ def load_data(self, data, data_label=None, do_link=True, show_in_viewer=True, **
 
                 kw = deepcopy(kwargs)
 
-                # This will only append index to data label if provided.
                 if data_label:
                     kw['data_label'] = data_label