data-menu: add subsets/data (#3263)

* import menu with add data to viewer support
* create subset support
* filter available data per-viewer
* basic test coverage

CHANGES.rst

@@ -4,7 +4,7 @@
 New Features
 ------------
 
-* New design for viewer legend. [#3220, #3254]
+* New design for viewer legend. [#3220, #3254, #3263]
 
 Cubeviz
 ^^^^^^^

jdaviz/app.py

@@ -151,7 +151,8 @@ def to_unit(self, data, cid, values, original_units, target_units):
                      'plugin-slider': 'components/plugin_slider.vue',
                      'plugin-color-picker': 'components/plugin_color_picker.vue',
                      'plugin-input-header': 'components/plugin_input_header.vue',
-                     'glue-state-sync-wrapper': 'components/glue_state_sync_wrapper.vue'}
+                     'glue-state-sync-wrapper': 'components/glue_state_sync_wrapper.vue',
+                     'data-menu-add-data': 'components/data_menu_add_data.vue'}
 
 _verbosity_levels = ('debug', 'info', 'warning', 'error')
 

jdaviz/components/data_menu_add_data.vue

@@ -0,0 +1,66 @@
+<template>
+  <v-menu
+    absolute
+    offset-y
+    left
+    v-if="dataset_items.length > 0 || subset_tools.length > 0"
+    >
+    <template v-slot:activator="{ on, attrs }">
+      <j-tooltip
+        v-if="dataset_items.length > 0 || subset_tools.length > 0"
+        tooltipcontent="Add data or subset to viewer"
+      >
+        <v-btn
+          icon
+          v-bind="attrs"
+          v-on="on"
+        >
+          <v-icon class="invert-if-dark">mdi-plus</v-icon>
+        </v-btn>
+      </j-tooltip>
+    </template>
+    <v-list dense style="width: 200px; max-height: 300px; overflow-y: auto;">
+      <v-subheader v-if="dataset_items.length > 0"><span>Load Data</span></v-subheader>
+      <v-list-item
+        v-for="data in dataset_items"
+      >
+        <v-list-item-content>
+          <j-tooltip tooltipcontent="add data to viewer">
+            <span
+              style="cursor: pointer; width: 100%"
+              @click="() => {$emit('add-data', data.label)}"
+            >
+              {{ data.label }}
+            </span>
+          </j-tooltip>
+        </v-list-item-content>
+      </v-list-item>
+      <v-subheader v-if="subset_tools.length > 0"><span>Create Subset</span></v-subheader>
+      <v-list-item
+        v-if="subset_tools.length > 0"
+      >
+        <v-list-item-content style="display: inline-block">
+          <j-tooltip
+            v-for="tool in subset_tools"
+            span_style="display: inline-block"
+            :tooltipcontent="'Create new '+tool.name+' subset'"
+          >
+            <v-btn 
+              icon
+              @click="() => {$emit('create-subset', tool.name)}"
+            >
+              <img :src="tool.img" width="20" class="invert-if-dark"/>
+            </v-btn>
+          </j-tooltip>
+        </v-list-item-content>
+      </v-list-item>
+    </v-list>
+  </v-menu>
+
+</template>
+
+<script>
+module.exports = {
+  props: ['dataset_items', 'subset_tools'],
+};
+</script>

jdaviz/components/toolbar_nested.py

@@ -156,6 +156,39 @@ def _on_change_v_model(self, event):
             if event['old'] is not None:
                 self.active_tool_id = event['old']
 
+    def _select_tool(self, tool_id, menu_ind):
+        for search_tool_id, info in self.tools_data.items():
+            if info['menu_ind'] == menu_ind and info['primary']:
+                prev_id = search_tool_id
+                prev_info = info
+                break
+        else:
+            raise ValueError("could not find previous selection")
+
+        if isinstance(self.tools[tool_id], CheckableTool):
+            # only switch to primary if its actually checkable, otherwise
+            # just activate once
+            self.tools_data = {
+                **self.tools_data,
+                prev_id: {
+                    **prev_info,
+                    'primary': False
+                },
+                tool_id: {
+                    **self.tools_data[tool_id],
+                    'primary': True
+                }
+            }
+
+        # and finally, set to be the active tool (this triggers _on_change_v_model which in turn
+        # triggers BasicJupyterToolbar._on_change_active_tool)
+        self.active_tool_id = tool_id
+
+    def select_tool(self, tool_id):
+        # find the previous primary tool in the same menu_ind
+        menu_ind = self.tools_data[tool_id]['menu_ind']
+        self._select_tool(tool_id, menu_ind)
+
     def add_tool(self, tool, menu_ind, has_suboptions=True, primary=False, visible=True):
         # NOTE: this method is essentially copied from glue-jupyter's BasicJupyterToolbar,
         # but we need extra values in the tools_data dictionary.  We could call super(),
@@ -183,29 +216,4 @@ def vue_select_primary(self, args):
         Activate the primary tool from a given menu index
         """
         menu_ind, tool_id = args
-        for search_tool_id, info in self.tools_data.items():
-            if info['menu_ind'] == menu_ind and info['primary']:
-                prev_id = search_tool_id
-                prev_info = info
-                break
-        else:
-            raise ValueError("could not find previous selection")
-
-        if isinstance(self.tools[tool_id], CheckableTool):
-            # only switch to primary if its actually checkable, otherwise
-            # just activate once
-            self.tools_data = {
-                **self.tools_data,
-                prev_id: {
-                    **prev_info,
-                    'primary': False
-                },
-                tool_id: {
-                    **self.tools_data[tool_id],
-                    'primary': True
-                }
-            }
-
-        # and finally, set to be the active tool (this triggers _on_change_v_model which in turn
-        # triggers BasicJupyterToolbar._on_change_active_tool)
-        self.active_tool_id = tool_id
+        self._select_tool(tool_id, menu_ind)

jdaviz/configs/cubeviz/plugins/viewers.py

@@ -39,6 +39,8 @@ def __init__(self, *args, **kwargs):
         # Hide axes by default
         self.state.show_axes = False
 
+        self.data_menu._obj.dataset.add_filter('is_cube_or_image')
+
     @property
     def _default_spectrum_viewer_reference_name(self):
         return self.jdaviz_helper._default_spectrum_viewer_reference_name

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

@@ -1,15 +1,25 @@
 from contextlib import contextmanager
 from traitlets import Bool, Dict, Unicode, List, observe
 
-from jdaviz.core.template_mixin import TemplateMixin, LayerSelectMixin
+from jdaviz.core.template_mixin import (TemplateMixin, LayerSelectMixin, DatasetSelectMixin)
 from jdaviz.core.user_api import UserApiWrapper
 from jdaviz.core.events import IconsUpdatedMessage, AddDataMessage
 from jdaviz.utils import cmap_samples, is_not_wcs_only
 
 __all__ = ['DataMenu']
 
 
-class DataMenu(TemplateMixin, LayerSelectMixin):
+SUBSET_TOOL_IDS = {'circle': 'bqplot:truecircle',
+                   'rectangle': 'bqplot:rectangle',
+                   'ellipse': 'bqplot:ellipse',
+                   'annulus': 'bqplot:circannulus',
+                   'xrange': 'bqplot:xrange',
+                   'yrange': 'bqplot:yrange'}
+
+SUBSET_NAMES = {v: k for k, v in SUBSET_TOOL_IDS.items()}
+
+
+class DataMenu(TemplateMixin, LayerSelectMixin, DatasetSelectMixin):
     """Viewer Data Menu
 
     Only the following attributes and methods are available through the
@@ -19,6 +29,8 @@ class DataMenu(TemplateMixin, LayerSelectMixin):
         actively selected layer(s)
     * :meth:`set_layer_visibility`
     * :meth:`toggle_layer_visibility`
+    * :meth:`create_subset`
+    * :meth:`add_data`
     """
     template_file = __file__, "data_menu.vue"
 
@@ -31,6 +43,7 @@ class DataMenu(TemplateMixin, LayerSelectMixin):
     visible_layers = Dict().tag(sync=True)  # read-only, set by viewer
 
     cmap_samples = Dict(cmap_samples).tag(sync=True)
+    subset_tools = List().tag(sync=True)
 
     dm_layer_selected = List().tag(sync=True)
 
@@ -47,18 +60,38 @@ def __init__(self, viewer, *args, **kwargs):
         self.layer.multiselect = True
         self.layer._default_mode = 'empty'
 
+        # we'll use a modified version of the dataset mixin to have a filtered
+        # list of data entries in the app that are not in the current viewer.
+        # changing the selection has no consequence.
+        def data_not_in_viewer(data):
+            return data.label not in self.layer.choices
+        self.dataset.filters = ['is_not_wcs_only', 'not_child_layer', data_not_in_viewer]
+
         # first attach callback to catch any updates to viewer/layer icons and then
         # set their initial state
         self.hub.subscribe(self, IconsUpdatedMessage, self._on_app_icons_updated)
         self.hub.subscribe(self, AddDataMessage, handler=lambda _: self._set_viewer_id())
         self.viewer_icons = dict(self.app.state.viewer_icons)
         self.layer_icons = dict(self.app.state.layer_icons)
 
+        # this currently assumes that toolbar.tools_data is set at init and does not change
+        # if we ever support dynamic tool registration, this will need to be updated
+        self.subset_tools = [{'id': k, 'img': v['img'], 'name': SUBSET_NAMES.get(k, k)}
+                             for k, v in self._viewer.toolbar.tools_data.items()
+                             if k in SUBSET_TOOL_IDS.values()]
+
     @property
     def user_api(self):
-        expose = ['layer', 'set_layer_visibility', 'toggle_layer_visibility']
+        expose = ['layer', 'set_layer_visibility', 'toggle_layer_visibility',
+                  'create_subset', 'add_data']
         return UserApiWrapper(self, expose=expose)
 
+    @observe('layer_items')
+    def _update_data_not_in_viewer(self, msg):
+        # changing the layers in the viewer needs to trigger an update to dataset_items
+        # through the set filters
+        self.dataset._on_data_changed()
+
     def _set_viewer_id(self):
         # viewer_ids are not populated on the viewer at init, so we'll keep checking and set
         # these the first time that they are available
@@ -162,3 +195,39 @@ def toggle_layer_visibility(self, layer_label):
 
     def vue_set_layer_visibility(self, info, *args):
         return self.set_layer_visibility(info.get('layer'), info.get('value'))  # pragma: no cover
+
+    def add_data(self, data_label):
+        """
+        Add a dataset to the viewer.
+
+        Parameters
+        ----------
+        data_label : str
+            The label of the dataset to add to the viewer.
+        """
+        if data_label not in self.dataset.choices:
+            raise ValueError(f"Data label '{data_label}' not able to be loaded into '{self.viewer_id}'.  Must be one of: {self.dataset.choices}")  # noqa
+        return self.app.add_data_to_viewer(self.viewer_id, data_label)
+
+    def vue_add_data_to_viewer(self, info, *args):
+        self.add_data(info.get('data_label'))  # pragma: no cover
+
+    def create_subset(self, subset_type):
+        """
+        Create a new subset in the viewer.  This sets the app-wide subset selection to 'Create New'
+        and selects the appropriate tool in this viewer's toolbar.
+
+        Parameters
+        ----------
+        subset_type : str
+            The type of subset to create.  Must be one of 'circle', 'rectangle', 'ellipse',
+            'annulus', 'xrange', or 'yrange'.
+        """
+        # clear previous selection, finalize subsets, temporarily sets default tool
+        self._viewer.toolbar.active_tool_id = None
+        # set toolbar to the selection, will also set app-wide subset selection to "Create New"
+        # NOTE: supports passing either the user-friendly name or the actual ID
+        self._viewer.toolbar.select_tool(SUBSET_TOOL_IDS.get(subset_type, subset_type))
+
+    def vue_create_subset(self, info, *args):
+        self.create_subset(info.get('subset_type'))  # pragma: no cover

jdaviz/configs/default/plugins/data_menu/data_menu.vue

@@ -63,7 +63,7 @@
                 icon
                 disabled
               >
-                <v-icon>mdi-format-vertical-align-center</v-icon>
+                <v-icon class="invert-if-dark">mdi-format-vertical-align-center</v-icon>
               </v-btn>
             </j-tooltip>
           </v-list-item-icon>
@@ -75,14 +75,13 @@
             </j-tooltip>
           </v-list-item-content>
           <v-list-item-action>
-            <j-tooltip tooltipcontent="Add data or subset (COMING SOON)">
-              <v-btn
-                icon
-                disabled
-              >
-                <v-icon>mdi-plus</v-icon>
-              </v-btn>
-            </j-tooltip>
+            <data-menu-add-data
+              :dataset_items="dataset_items"
+              :subset_tools="subset_tools"
+              @add-data="(data_label) => {add_data_to_viewer({data_label: data_label})}"
+              @create-subset="(subset_type) => {create_subset({subset_type: subset_type}); data_menu_open = false}"
+            >
+            </data-menu-add-data>
           </v-list-item-action>
         </v-list-item>
         <v-list-item-group
@@ -144,7 +143,7 @@
                 icon
                 disabled
               >
-                <v-icon>mdi-delete</v-icon>
+                <v-icon class="invert-if-dark">mdi-delete</v-icon>
               </v-btn>
             </j-tooltip>
             <j-tooltip
@@ -156,7 +155,7 @@
                 icon
                 disabled
                 >
-                <v-icon>mdi-label</v-icon>
+                <v-icon class="invert-if-dark">mdi-label</v-icon>
               </v-btn>
             </j-tooltip>
             <j-tooltip
@@ -166,6 +165,7 @@
               <v-btn
                 text
                 disabled
+                class="invert-if-dark"
               >
                 Edit Subset
               </v-btn>
@@ -271,6 +271,7 @@
   .dm-footer > .v-list-item__icon, .dm-footer > .v-list-item__content, .dm-footer > .v-list-item__action {
     filter: invert(1);
   }
+
   .dm-header.v-btn--disabled  .v-icon {
     color: green !important;
   }

jdaviz/configs/default/tests/test_data_menu.py

@@ -66,3 +66,28 @@ def test_data_menu_selection(specviz_helper, spectrum1d):
     assert len(dm._obj.layer_items) == 1
     assert dm._obj.dm_layer_selected == [0]
     assert dm.layer.selected == ['test']
+
+
+def test_data_menu_add_data(imviz_helper):
+    for i in range(3):
+        imviz_helper.load_data(np.zeros((2, 2)) + i, data_label=f'image_{i}', show_in_viewer=False)
+
+    dm = imviz_helper.viewers['imviz-0']._obj.data_menu
+    assert len(dm._obj.layer_items) == 0
+    assert len(dm.layer.choices) == 0
+    assert len(dm._obj.dataset.choices) == 3
+
+    dm.add_data('image_0')
+    assert dm.layer.choices == ['image_0']
+    assert len(dm._obj.dataset.choices) == 2
+
+
+def test_data_menu_create_subset(imviz_helper):
+    imviz_helper.load_data(np.zeros((2, 2)), data_label='image', show_in_viewer=True)
+
+    dm = imviz_helper.viewers['imviz-0']._obj.data_menu
+    assert imviz_helper.app.session.edit_subset_mode.edit_subset == []
+
+    dm.create_subset('circle')
+    assert imviz_helper.app.session.edit_subset_mode.edit_subset == []
+    assert imviz_helper.viewers['imviz-0']._obj.toolbar.active_tool_id == 'bqplot:truecircle'

jdaviz/configs/imviz/plugins/viewers.py

@@ -63,6 +63,8 @@ def __init__(self, *args, **kwargs):
         # intensive.
         self.state.image_external_padding = 0.5
 
+        self.data_menu._obj.dataset.add_filter('is_image')
+
     def on_mouse_or_key_event(self, data):
         active_image_layer = self.active_image_layer
         if active_image_layer is None: