Fix some bugs in the Jupyter tool.

glue_ar/jupyter/export_tool.py

@@ -1,11 +1,16 @@
 from os import getcwd
-from os.path import exists
+from os.path import exists, splitext
+from typing import Dict, Iterable, Tuple
 
 from glue.config import viewer_tool
+from glue.core.state_objects import State
+from glue.viewers.common.state import LayerState
 from glue.viewers.common.tool import Tool
+from glue_vispy_viewers.volume.qt.volume_viewer import VispyVolumeViewerMixin
 
-from glue_ar.common.export import export_to_ar
-from glue_ar.utils import AR_ICON
+from glue_ar.common.export import export_viewer
+from glue_ar.common.export_options import ar_layer_export
+from glue_ar.utils import AR_ICON, export_label_for_layer, xyz_bounds
 
 import ipyvuetify as v  # noqa
 from ipywidgets import HBox, Layout  # noqa
@@ -45,6 +50,7 @@ def on_ok_click(button, event, data):
 
         def on_close_click(button, event, data):
             self.viewer.output_widget.clear_output()
+            dialog.close()
 
         def on_selected_change(chooser):
             ok_btn.disabled = not bool(chooser.selected_filename)
@@ -88,5 +94,33 @@ def on_no_click(button, event, data):
             self.save_figure(filepath)
             self.viewer.output_widget.clear_output()
 
+    def _state_dictionary(self,
+                          layers: Iterable[LayerState],
+                          filetype: str) -> Dict[str, Tuple[str, State]]:
+        state_dict = {}
+        for layer in layers:
+            label = export_label_for_layer(layer)
+            layer_state_cls = type(layer.state)
+            states = ar_layer_export.export_state_classes(layer_state_cls)
+            method_names = ar_layer_export.method_names(layer_state_cls, filetype)
+            method = method_names[0]
+            state_cls = next(t[1] for t in states if t[0] == method)
+            state = state_cls()
+            state_dict[label] = (method, state)
+
+        return state_dict
+
     def save_figure(self, filepath):
-        export_to_ar(self.viewer, filepath, state_dict={}, compression="draco")
+        bounds = xyz_bounds(self.viewer.state, with_resolution=isinstance(self.viewer, VispyVolumeViewerMixin))
+        layers = [layer for layer in self.viewer.layers if layer.enabled and layer.state.visible]
+        filetype = splitext(filepath)[1][1:]
+        state_dict = self._state_dictionary(layers, filetype)
+        export_viewer(viewer_state=self.viewer.state,
+                      layer_states=[layer.state for layer in self.viewer.layers
+                                    if layer.enabled and layer.state.visible],
+                      bounds=bounds,
+                      state_dictionary=state_dict,
+                      filepath=filepath)
+
+
+