data-menu: handle viewer rename properly

spacetelescope · Jan 8, 2025 · 665ae6b · 665ae6b
1 parent 19ddfac
commit 665ae6b
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 1 deletion.
diff --git a/jdaviz/configs/default/plugins/data_menu/data_menu.py b/jdaviz/configs/default/plugins/data_menu/data_menu.py
@@ -6,7 +6,8 @@
 from jdaviz.core.template_mixin import (TemplateMixin, LayerSelect,
                                         LayerSelectMixin, DatasetSelectMixin)
 from jdaviz.core.user_api import UserApiWrapper
-from jdaviz.core.events import IconsUpdatedMessage, AddDataMessage, ChangeRefDataMessage
+from jdaviz.core.events import (IconsUpdatedMessage, AddDataMessage,
+                                ChangeRefDataMessage, ViewerRenamedMessage)
 from jdaviz.utils import cmap_samples, is_not_wcs_only
 
 from glue.core.edit_subset_mode import (AndMode, AndNotMode, OrMode,
@@ -132,6 +133,7 @@ def data_not_in_viewer(data):
         self.hub.subscribe(self, IconsUpdatedMessage, self._on_app_icons_updated)
         self.hub.subscribe(self, AddDataMessage, handler=lambda _: self._set_viewer_id())
         self.hub.subscribe(self, ChangeRefDataMessage, handler=self._on_refdata_change)
+        self.hub.subscribe(self, ViewerRenamedMessage, handler=self._on_viewer_renamed_message)
         self.viewer_icons = dict(self.app.state.viewer_icons)
         self.layer_icons = dict(self.app.state.layer_icons)
 
@@ -210,6 +212,11 @@ def _on_refdata_change(self, msg):
             with self.during_select_sync():
                 self.orientation.selected = str(self._viewer.state.reference_data.label)
 
+    def _on_viewer_renamed_message(self, msg):
+        if self.viewer_reference == msg.old_viewer_ref:
+            self.viewer_reference = msg.new_viewer_ref
+            self._set_viewer_id()
+
     @observe('orientation_layer_selected')
     def _orientation_layer_selected_changed(self, event={}):
         if not hasattr(self, 'orientation'):

diff --git a/jdaviz/core/template_mixin.py b/jdaviz/core/template_mixin.py
@@ -1537,6 +1537,8 @@ def __init__(self, plugin, items, selected, viewer,
                            handler=lambda _: self._update_items())
         self.hub.subscribe(self, SubsetRenameMessage,
                            handler=self._on_subset_renamed)
+        self.hub.subscribe(self, ViewerRenamedMessage,
+                           self._on_viewer_renamed_message)
 
         self.sort_by = sort_by
         self.app.state.add_callback('layer_icons', self._update_items)
@@ -1575,6 +1577,14 @@ def filter_has_children(data):
 
             self.add_filter(filter_has_children)
 
+    def _on_viewer_renamed_message(self, msg):
+        if isinstance(self.viewer, list):
+            for i, viewer in self.viewer:
+                if viewer == msg.old_viewer_ref:
+                    self.viewer[i] = msg.new_viewer_ref
+        elif self.viewer == msg.old_viewer_ref:
+            self.viewer = msg.new_viewer_ref
+
     def _get_viewer(self, viewer):
         # newer will likely be the viewer name in most cases, but viewer id in the case
         # of additional viewers in imviz.