From 12acaf01c30c65a499be68ef234bdc9d9a9ed73b Mon Sep 17 00:00:00 2001
From: Kyle Conroy <kyleconroy@gmail.com>
Date: Wed, 8 Nov 2023 13:03:27 -0500
Subject: [PATCH] LayerSelect: fix wcs-only filtering

---
 .../plugins/links_control/links_control.py    |  2 +-
 jdaviz/core/template_mixin.py                 | 36 ++++++++-----------
 2 files changed, 15 insertions(+), 23 deletions(-)

diff --git a/jdaviz/configs/imviz/plugins/links_control/links_control.py b/jdaviz/configs/imviz/plugins/links_control/links_control.py
index 5f1668f16d..8ef14e29b4 100644
--- a/jdaviz/configs/imviz/plugins/links_control/links_control.py
+++ b/jdaviz/configs/imviz/plugins/links_control/links_control.py
@@ -98,7 +98,7 @@ def __init__(self, *args, **kwargs):
 
         self.layer = LayerSelect(
             self, 'layer_items', 'layer_selected', 'viewer_selected',
-            'multiselect', only_wcs_layers=True
+            'multiselect', filters=['is_wcs_only']
         )
         self.orientation_layer_label = AutoTextField(
             self, 'new_layer_label', 'new_layer_label_default', 'new_layer_label_auto', None
diff --git a/jdaviz/core/template_mixin.py b/jdaviz/core/template_mixin.py
index f953bf3615..563abffb4f 100644
--- a/jdaviz/core/template_mixin.py
+++ b/jdaviz/core/template_mixin.py
@@ -1202,7 +1202,7 @@ def __init__(self, plugin, items, selected, viewer,
                  multiselect=None,
                  default_text=None, manual_options=[],
                  default_mode='first',
-                 only_wcs_layers=False):
+                 filters=['is_not_wcs_only']):
         """
         Parameters
         ----------
@@ -1223,7 +1223,6 @@ def __init__(self, plugin, items, selected, viewer,
             ``default`` text is provided but not in ``manual_options`` it will still be included as
             the first item in the list.
         """
-
         super().__init__(plugin,
                          items=items,
                          selected=selected,
@@ -1233,8 +1232,6 @@ def __init__(self, plugin, items, selected, viewer,
                          manual_options=manual_options,
                          default_mode=default_mode)
 
-        self.only_wcs_layers = only_wcs_layers
-
         self.hub.subscribe(self, AddDataMessage,
                            handler=lambda _: self._on_layers_changed())
         self.hub.subscribe(self, RemoveDataMessage,
@@ -1249,7 +1246,7 @@ def __init__(self, plugin, items, selected, viewer,
 
         self.app.state.add_callback('layer_icons', lambda _: self._on_layers_changed())
         self.add_observe(viewer, self._on_viewer_changed)
-        self._on_layers_changed()
+        self.filters = filters
 
     def _get_viewer(self, viewer):
         # newer will likely be the viewer name in most cases, but viewer id in the case
@@ -1276,6 +1273,15 @@ def _on_viewer_changed(self, msg=None):
             self._clear_cache()
             self._on_layers_changed()
 
+    def _is_valid_item(self, layer):
+        def is_wcs_only(layer):
+            return hasattr(layer.layer, 'meta') and layer.layer.meta.get('_WCS_ONLY', False)
+
+        def is_not_wcs_only(layer):
+            return not is_wcs_only(layer)
+
+        return super()._is_valid_item(layer, locals())
+
     @observe('filters')
     def _on_layers_changed(self, msg=None):
         # NOTE: _on_layers_changed is passed without a msg object during init
@@ -1289,23 +1295,9 @@ def _on_layers_changed(self, msg=None):
 
         # use getattr so the super() call above doesn't try to access the attr before
         # it is initialized:
-        if not getattr(self, 'only_wcs_layers', False):
-            layers = [
-                layer for viewer in viewers
-                for layer in getattr(viewer, 'layers', [])
-                # don't include WCS-only layers unless asked:
-                if (
-                    not hasattr(layer.layer, 'meta') or
-                    (not layer.layer.meta.get('_WCS_ONLY', False))
-                )
-            ]
-        else:
-            layers = [
-                layer for viewer in viewers
-                for layer in getattr(viewer, 'layers', [])
-                # only include WCS-only layers:
-                if not hasattr(layer.layer, 'meta') or layer.layer.meta.get('_WCS_ONLY', False)
-            ]
+        layers = [layer for viewer in viewers for layer in getattr(viewer, 'layers', [])
+                  if self._is_valid_item(layer)]
+
         # remove duplicates - NOTE: by doing this, any color-mismatch between layers with the
         # same name in different viewers will be randomly assigned within plot_options
         # based on which was found _first.