Add data export functionality (spacetelescope#2770)

* Add data export functionality * Fix code style * Add a change log * Select dataset format from a dropdown * Add a test to export data * codestyle * Test to unable export data when unsupported units * Update jdaviz/configs/default/plugins/export/export.py Co-authored-by: Kyle Conroy <[email protected]> * Disable export for NDdata obj * Update jdaviz/configs/default/plugins/export/export.py Co-authored-by: Kyle Conroy <[email protected]> * Fix code style * Fix test * Disable export for non plugin generated data * fix codestyle * Fix test export data * Check invalid msg in test * Codestyle --------- Co-authored-by: Kyle Conroy <[email protected]>
rosteen · Apr 2, 2024 · c31424d · c31424d
1 parent 77c9d44
commit c31424d
Show file tree

Hide file tree

Showing 4 changed files with 89 additions and 7 deletions.
diff --git a/CHANGES.rst b/CHANGES.rst
@@ -15,7 +15,7 @@ New Features
 
 - "Export Plot" plugin is now replaced with the more general "Export" plugin. [#2722]
 
-- "Export" plugin supports exporting plugin tables and non-composite spatial subsets.[#2755, #2760, #2772]
+- "Export" plugin supports exporting plugin tables, data and non-composite spatial subsets.[#2755, #2760, #2772, #2770]
 
 - Opening a plugin in the tray (from the API or the toolbar buttons) now scrolls to that plugin.
   [#2768]

diff --git a/jdaviz/configs/default/plugins/export/export.py b/jdaviz/configs/default/plugins/export/export.py
@@ -13,6 +13,8 @@
 
 from jdaviz.core.events import AddDataMessage, SnackbarMessage
 from jdaviz.core.user_api import PluginUserApi
+from specutils import Spectrum1D
+from astropy.nddata import CCDData
 
 try:
     import cv2
@@ -64,10 +66,14 @@ class Export(PluginTemplateMixin, ViewerSelectMixin, SubsetSelectMixin,
     subset_format_items = List().tag(sync=True)
     subset_format_selected = Unicode().tag(sync=True)
 
+    dataset_format_items = List().tag(sync=True)
+    dataset_format_selected = Unicode().tag(sync=True)
+
     filename = Unicode().tag(sync=True)
 
     # if selected subset is spectral or composite, display message and disable export
     subset_invalid_msg = Unicode().tag(sync=True)
+    data_invalid_msg = Unicode().tag(sync=True)
 
     # We currently disable exporting spectrum-viewer in Cubeviz
     viewer_invalid_msg = Unicode().tag(sync=True)
@@ -117,6 +123,12 @@ def __init__(self, *args, **kwargs):
                                                    selected='subset_format_selected',
                                                    manual_options=subset_format_options)
 
+        dataset_format_options = ['fits']
+        self.dataset_format = SelectPluginComponent(self,
+                                                    items='dataset_format_items',
+                                                    selected='dataset_format_selected',
+                                                    manual_options=dataset_format_options)
+
         # default selection:
         self.dataset._default_mode = 'empty'
         self.table._default_mode = 'empty'
@@ -141,12 +153,11 @@ def user_api(self):
         # TODO: expose export method once API is finalized
 
         expose = ['viewer', 'viewer_format',
+                  'dataset', 'dataset_format',
                   'subset', 'subset_format',
                   'table', 'table_format',
                   'filename', 'export']
 
-        if self.dev_dataset_support:
-            expose += ['dataset']
         if self.dev_plot_support:
             expose += ['plot']
         if self.dev_multi_support:
@@ -190,6 +201,8 @@ def _sync_singleselect(self, event):
                 setattr(self, attr, '')
             if attr == 'subset_selected':
                 self._set_subset_not_supported_msg()
+            if attr == 'dataset_selected':
+                self._set_dataset_not_supported_msg()
             elif self.config == "cubeviz" and attr == "viewer_selected":
                 self._disable_viewer_format_combo(event)
 
@@ -221,6 +234,17 @@ def _set_subset_not_supported_msg(self, msg=None):
         else:  # no subset selected (can be '' instead of None if previous selection made)
             self.subset_invalid_msg = ''
 
+    def _set_dataset_not_supported_msg(self, msg=None):
+        if self.dataset.selected_obj is not None:
+            if self.dataset.selected_obj.meta.get("Plugin", None) is None:
+                self.data_invalid_msg = "Data export is only available for plugin generated data."
+            elif not isinstance(self.dataset.selected_obj, (Spectrum1D, CCDData)):
+                self.data_invalid_msg = "Export is not implemented for this type of data"
+            else:
+                self.data_invalid_msg = ''
+        else:
+            self.data_invalid_msg = ''
+
     @with_spinner()
     def export(self, filename=None, show_dialog=None):
         """
@@ -231,8 +255,6 @@ def export(self, filename=None, show_dialog=None):
         filename : str, optional
             If not provided, plugin value will be used.
         """
-        if self.dataset.selected is not None and len(self.dataset.selected):
-            raise NotImplementedError("dataset export not yet supported")
 
         if self.plot.selected is not None and len(self.plot.selected):
             raise NotImplementedError("plot export not yet supported")
@@ -273,6 +295,13 @@ def export(self, filename=None, show_dialog=None):
                 raise NotImplementedError(f'Subset can not be exported - {self.subset_invalid_msg}')
             self.save_subset_as_region(selected_subset_label, filename)
 
+        elif len(self.dataset.selected):
+            filetype = self.dataset_format.selected
+            if not filename.endswith(filetype):
+                filename += f".{filetype}"
+            if self.data_invalid_msg != "":
+                raise NotImplementedError(f"Data can not be exported - {self.data_invalid_msg}")
+            self.dataset.selected_obj.write(Path(filename), overwrite=True)
         else:
             raise ValueError("nothing selected for export")
         return filename

diff --git a/jdaviz/configs/default/plugins/export/export.vue b/jdaviz/configs/default/plugins/export/export.vue
@@ -84,7 +84,7 @@
       </div>
     </div>
 
-    <div v-if="dev_dataset_support && dataset_items.length > 0">
+    <div v-if="dataset_items.length > 0">
       <j-plugin-section-header style="margin-top: 12px">Data</j-plugin-section-header>
       <plugin-inline-select
         :items="dataset_items"
@@ -94,6 +94,23 @@
       >
       </plugin-inline-select>
     </div>
+    <v-row v-if="data_invalid_msg.length > 0">
+      <span class="v-messages v-messages__message text--secondary" style="color: red !important">
+                {{data_invalid_msg}}
+      </span>
+    </v-row>
+    <v-row v-if="dataset_selected" class="row-min-bottom-padding">
+        <v-select
+          :menu-props="{ left: true }"
+          attach
+          v-model="dataset_format_selected"
+          :items="dataset_format_items.map(i => i.label)"
+          label="Format"
+          hint="Format for exporting datasets."
+          persistent-hint
+        >
+        </v-select>
+      </v-row>
 
     <div v-if="subset_items.length > 0">
       <j-plugin-section-header style="margin-top: 12px">Subsets</j-plugin-section-header>
@@ -191,7 +208,7 @@
         :spinner="spinner"
         :disabled="filename.length === 0 ||
                    movie_recording ||
-                   subset_invalid_msg.length > 0 ||
+                   subset_invalid_msg.length > 0 || data_invalid_msg.length > 0 ||
                    viewer_invalid_msg.length > 0 ||
                    (viewer_selected.length > 0 && viewer_format_selected == 'mp4' && !movie_enabled)"
       >

diff --git a/jdaviz/configs/default/plugins/export/tests/test_export.py b/jdaviz/configs/default/plugins/export/tests/test_export.py
@@ -183,3 +183,39 @@ def test_basic_export_subsets_cubeviz(self, tmp_path, cubeviz_helper, spectral_c
         with pytest.raises(NotImplementedError,
                            match='Subset can not be exported - Export for composite subsets not supported.'):  # noqa
             export_plugin.export()
+
+
+def test_export_data(cubeviz_helper, spectrum1d_cube):
+    cubeviz_helper.load_data(spectrum1d_cube, data_label="test")
+    export_plugin = cubeviz_helper.plugins["Export"]._obj
+    export_plugin.dataset_selected = 'test[FLUX]'
+    assert export_plugin.dataset_format.selected == 'fits'
+    assert 'test[FLUX]' in cubeviz_helper.app.data_collection.labels
+    with pytest.raises(NotImplementedError,
+                       match='Data can not be exported'):
+        export_plugin.export()
+    assert export_plugin.data_invalid_msg == 'Data export is only available for plugin generated data.'   # noqa
+
+
+def test_disable_export_for_non_plugin_generated_data(cubeviz_helper, spectrum1d_cube):
+    cubeviz_helper.load_data(spectrum1d_cube, data_label='test')
+    mm = cubeviz_helper.plugins["Moment Maps"]
+    mm.dataset = 'test[FLUX]'
+    mm._obj.dataset_selected = 'test[FLUX]'
+    mm.n_moment = 0
+    mm.calculate_moment()
+    assert mm._obj.results_label == 'moment 0'
+    cubeviz_helper.app.add_data_to_viewer(
+        cubeviz_helper._default_flux_viewer_reference_name, 'moment 0'
+    )
+    ep = cubeviz_helper.plugins["Export"]._obj
+    ep.dataset_selected = 'test[FLUX]'
+    with pytest.raises(NotImplementedError,
+                       match='Data can not be exported'):
+        ep.export()
+    assert ep.data_invalid_msg == 'Data export is only available for plugin generated data.'
+
+    ep.dataset_selected = 'moment 0'
+    ep.export()
+    assert os.path.isfile("cubeviz_export.fits")
+    assert ep.data_invalid_msg == ''