feat: export(block=True) will block next cell execution until export finishes

Author: maartenbreddels

jdaviz/async_utils.py

@@ -0,0 +1,149 @@
+"""Utility function for working with asyncio and widgets."""
+import asyncio
+import time
+import logging
+import threading
+
+from typing import Callable
+import ipywidgets
+from asyncio import Queue
+import queue
+import solara
+import solara.lab
+
+logger = logging.getLogger(__name__)
+
+# instead of a single global variable, this makes it kernel
+# scoped in solara
+_serial_task_queue = solara.lab.computed(Queue)
+
+
+# analogous to asyncio.create_task, except, they will be called in sequence
+def create_serial_task(coro):
+    serial_task_run_task.get()  # ensure task runner is running
+    _serial_task_queue.value.put_nowait(coro)
+
+
+async def serial_task_run():
+    logger.debug("serial_task_run: starting")
+    while True:
+        try:
+            logger.debug("serial_task_run: getting task from queue")
+            task = await _serial_task_queue.value.get()
+            logger.debug("serial_task_run: got task from queue, running it")
+            await task
+        except Exception:
+            logger.exception("Task failed")
+        finally:
+            _serial_task_queue.value.task_done()
+
+
+@solara.lab.computed
+def serial_task_run_task():
+    event_loop_queue = queue.Queue()
+
+    def runner():
+        try:
+            event_loop = asyncio.new_event_loop()
+            event_loop_queue.put(event_loop)
+
+            async def tick():
+                while True:
+                    await asyncio.sleep(0.1)
+                    print(".", end="", flush=True)
+
+            # not sure why we need this
+            event_loop.run_until_complete(tick())
+            # event_loop.run_until_complete(serial_task_run())
+            # but this doesn't work?
+            # event_loop.run_forever() doesn't work?
+        except Exception:
+            logger.exception("Task running thread failed")
+
+    thread = threading.Thread(target=runner, daemon=True)
+    thread.start()
+    event_loop = event_loop_queue.get()
+    event_loop.create_task(serial_task_run())
+    logger.debug(f"serial_task_run_task: created event loop: {event_loop}")
+    return event_loop
+
+
+def wait_for_change(widget: ipywidgets.Widget, trait_name: str):
+    # based on https://ipywidgets.readthedocs.io/en/latest/examples/Widget%20Asynchronous.html
+    future = asyncio.Future()
+
+    def getvalue(change):
+        logger.debug(
+            f"got new value for {trait_name} for widget {widget} new value {change.new}"
+        )
+        future.set_result(change.new)
+        widget.unobserve(getvalue, trait_name)
+
+    logger.debug(f"observing {trait_name} for widget {type(widget)}")
+    widget.observe(getvalue, trait_name)
+    return future
+
+
+async def queue_screenshot_async(widget, method, *args, **kwargs):
+    logger.debug("queue_screenshot_async: starting")
+    # queue with just 1 item, the callback data
+    data_queue = Queue()
+
+    def callback_wrapper(data):
+        logger.debug(
+            "queue_screenshot_async: callback_wrapper: putting result in queue"
+        )
+        try:
+            data_queue.put_nowait(data)
+        except Exception:
+            logger.exception("Failed to put data in queue")
+
+    # calls Figure.get_svg_data or something similar
+    method(callback_wrapper, *args, **kwargs)
+    return await data_queue.get()
+
+
+def queue_screenshot_sync(
+    widget, method, callback, on_timeout, timeout=5, *args, **kwargs
+):
+    # queue with just 1 item, the callback data
+    data_queue = Queue()
+
+    def callback_wrapper(data):
+        data_queue.put_nowait(data)
+
+    async def execute():
+        # calls Figure.get_svg_data or something similar
+        method(callback_wrapper, *args, **kwargs)
+        try:
+            data = await asyncio.wait_for(data_queue.get(), timeout=timeout)
+        except TimeoutError:
+            on_timeout()
+        callback(data)
+
+    create_serial_task(execute())
+
+
+def run_kernel_events_blocking_until(
+    condition: Callable[[], bool], timeout: float = 5, sleep: float = 0.1
+):
+    """Executes kernel events while the condition is true or the timeout is reached.
+
+    Used to block in the notebook while we wait for a widget result.
+    """
+    try:
+        from jupyter_ui_poll import ui_events
+    except ModuleNotFoundError:
+        raise ModuleNotFoundError(
+            "jupyter_ui_poll is not installed. Please install it with `pip install jupyter_ui_poll`"
+        )
+
+    start_time = time.time()
+    with ui_events() as poll:
+        while condition():
+            poll(10)  # process up to 10 UI events per iteration
+            time.sleep(sleep)
+            if time.time() - start_time > timeout:
+                raise TimeoutError(
+                    f"Timeout waiting for condition to be true after {timeout} seconds"
+                )

jdaviz/configs/default/plugins/export/export.py

@@ -1,8 +1,15 @@
+import logging
 import os
 import time
 from pathlib import Path
 import threading
 
+import asyncio
+import ipywidgets as widgets
+
+from jdaviz.async_utils import (create_serial_task, queue_screenshot_async,
+                                run_kernel_events_blocking_until,
+                                serial_task_run_task, wait_for_change)
 from astropy import units as u
 from astropy.nddata import CCDData
 from glue.core.message import SubsetCreateMessage, SubsetDeleteMessage, SubsetUpdateMessage
@@ -31,6 +38,7 @@
     HAS_OPENCV = True
 
 __all__ = ['Export']
+logger = logging.getLogger(__name__)
 
 
 @tray_registry('export', label="Export",
@@ -118,6 +126,7 @@ class Export(PluginTemplateMixin, ViewerSelectMixin, SubsetSelectMixin,
     # This is a temporary measure to allow server-installations to disable saving server-side until
     # saving client-side is supported for all exports.
     serverside_enabled = Bool(True).tag(sync=True)
+    _busy_doing_export = Bool(False).tag(sync=True)
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
@@ -443,7 +452,7 @@ def _normalize_filename(self, filename=None, filetype=None, overwrite=False, def
 
     @with_spinner()
     def export(self, filename=None, show_dialog=None, overwrite=False,
-               raise_error_for_overwrite=True):
+               raise_error_for_overwrite=True, block=False):
         """
         Export selected item(s)
 
@@ -462,6 +471,11 @@ def export(self, filename=None, show_dialog=None, overwrite=False,
             If `True`, raise exception when ``overwrite=False`` but
             output file already exists. Otherwise, a message will be sent
             to application snackbar instead.
+
+        block : bool
+            If `True`, block until the export is complete, this is useful in
+            a notebook context to ensure the export is complete before the
+            next export is started.
         """
         if self.multiselect:
             raise NotImplementedError("batch export not yet supported")
@@ -508,7 +522,7 @@ def export(self, filename=None, show_dialog=None, overwrite=False,
             else:
                 self.save_figure(viewer, filename, filetype, show_dialog=show_dialog,
                                  width=f"{self.image_width}px" if self.image_custom_size else None,
-                                 height=f"{self.image_height}px" if self.image_custom_size else None)  # noqa
+                                 height=f"{self.image_height}px" if self.image_custom_size else None, block=block)  # noqa
 
             # restore marks to their original state
             for restore, mark in zip(restores, viewer.figure.marks):
@@ -534,7 +548,7 @@ def export(self, filename=None, show_dialog=None, overwrite=False,
                     raise FileExistsError(f"{filename} exists but overwrite={overwrite}")
                 return
 
-            self.save_figure(plot, filename, filetype, show_dialog=show_dialog)
+            self.save_figure(plot, filename, filetype, show_dialog=show_dialog, block=block)
 
         elif len(self.plugin_table.selected):
             filetype = self.plugin_table_format.selected
@@ -609,17 +623,33 @@ def vue_overwrite_from_ui(self, *args, **kwargs):
             self.overwrite_warn = False
 
     def save_figure(self, viewer, filename=None, filetype="png", show_dialog=False,
-                    width=None, height=None):
+                    width=None, height=None, block=True):
         if filename is None:
             filename = self.filename_default
 
-        # viewers in plugins will have viewer.app, other viewers have viewer.jdaviz_app
-        if hasattr(viewer, 'jdaviz_app'):
-            app = viewer.jdaviz_app
-        else:
-            app = viewer.app
+        if self._busy_doing_export:
+            raise ValueError("Saving figure is still in progress. Use `save_figure(..., block=True)` or export(..., block=True) to make sure the previous export is complete") # noqa
+        self._busy_doing_export = True
 
-        def on_img_received(data):
+        async def save_figure_task():
+            try:
+                await self._save_figure_async(viewer, filename, filetype, show_dialog, width, height)
+            finally:
+                self._busy_doing_export = False
+        if block:
+            event_loop = serial_task_run_task.get()
+            logger.warning(f"event loop: {event_loop}, now creating task")
+            event_loop.create_task(save_figure_task())
+            run_kernel_events_blocking_until(lambda: not self._busy_doing_export)
+        else:
+            task = asyncio.create_task(save_figure_task())
+            create_serial_task(task)
+            return task
+
+    async def _save_figure_async(self, viewer, filename, filetype, show_dialog, width, height):
+        # Things become a bit more easy to reason about using async/await instead of callbacks
+        # So this internal method uses async/await instead of callbacks.
+        def save_to_file(data):
             try:
                 with filename.open(mode='bw') as f:
                     f.write(data)
@@ -632,17 +662,15 @@ def on_img_received(data):
                     f"{self.viewer.selected} exported to {str(filename)}",
                     sender=self, color="success"))
 
-        def get_png(figure):
-            if figure._upload_png_callback is not None:
-                raise ValueError("previous png export is still in progress. Wait to complete before making another call to save_figure")  # noqa: E501 # pragma: no cover
-
-            figure.get_png_data(on_img_received)
+        # viewers in plugins will have viewer.app, other viewers have viewer.jdaviz_app
+        if hasattr(viewer, 'jdaviz_app'):
+            app = viewer.jdaviz_app
+        else:
+            app = viewer.app
 
         if (width is not None or height is not None):
             assert width is not None and height is not None, \
                 "Both width and height must be provided"
-            import ipywidgets as widgets
-            from typing import Callable
 
             def _show_hidden(widget: widgets.Widget, width: str, height: str):
                 import ipyvuetify as v
@@ -658,42 +686,48 @@ def _show_hidden(widget: widgets.Widget, width: str, height: str):
                 # TODO: we might want to remove it from the DOM
                 app.invisible_children = [*app.invisible_children, wrapper_widget]
 
-            def _widget_after_first_display(widget: widgets.Widget, callback: Callable):
+            def _widget_after_first_display(widget: widgets.Widget):
                 if widget._view_count is None:
                     widget._view_count = 0
-                called_callback = False
-
-                def view_count_changed(change):
-                    nonlocal called_callback
-                    if change["new"] == 1 and not called_callback:
-                        called_callback = True
-                        callback()
-                widget.observe(view_count_changed, "_view_count")
+                logger.debug(f"waiting for view count to change for widget {type(widget)}")
+                return wait_for_change(widget, "_view_count")
 
             cloned_viewer = viewer._clone_viewer_outside_app()
             # make sure we will the size of our container which defines the
             # size of the figure
             cloned_viewer.figure.layout.width = "100%"
             cloned_viewer.figure.layout.height = "100%"
 
-            def on_figure_displayed():
-                # we need a bit of a delay to ensure the figure is fully displayed
-                # maybe this can be fixed on the bqplot side in the future
-                def wait_in_other_thread():
-                    import time
-                    time.sleep(0.2)
-                    get_png(cloned_viewer.figure)
-                # wait in other thread to avoid blocking the main thread (widgets can update)
-                threading.Thread(target=wait_in_other_thread).start()
-            _widget_after_first_display(cloned_viewer.figure, on_figure_displayed)
+            logger.debug("calling _widget_after_first_display for widget")
+            display_future = _widget_after_first_display(cloned_viewer.figure)
+            logger.debug(f"calling _show_hidden for widget {display_future}")
             _show_hidden(cloned_viewer.figure, width, height)
-        elif filetype == 'png':
-            # NOTE: get_png already check if _upload_png_callback is not None
-            get_png(viewer.figure)
-        elif filetype == 'svg':
+            logger.debug("waiting for display future")
+            await display_future
+            logger.debug("display future done")
+            await asyncio.sleep(0.2)
+            logger.debug("sleeping done")
+            if cloned_viewer.figure._upload_png_callback is not None:
+                raise ValueError("previous svg export is still in progress. Wait to complete before making another call to save_figure")
+            if cloned_viewer.figure._upload_svg_callback is not None:
+                raise ValueError("previous svg export is still in progress. Wait to complete before making another call to save_figure")
+            logger.debug("queueing screenshot")
+            get_image_data_method = cloned_viewer.figure.get_svg_data if filetype == 'svg' else cloned_viewer.figure.get_png_data
+            data = await queue_screenshot_async(cloned_viewer.figure, get_image_data_method)
+            logger.debug("got data, saving to file {filename}")
+            save_to_file(data)
+            logger.debug("saved to file {filename}")
+        elif filetype in ['png', 'svg']:
+            if viewer.figure._upload_png_callback is not None:
+                raise ValueError("previous png export is still in progress. Wait to complete before making another call to save_figure")
             if viewer.figure._upload_svg_callback is not None:
-                raise ValueError("previous svg export is still in progress. Wait to complete before making another call to save_figure") # noqa
-            viewer.figure.get_svg_data(on_img_received)
+                raise ValueError("previous svg export is still in progress. Wait to complete before making another call to save_figure")
+            get_image_data_method = viewer.figure.get_svg_data if filetype == 'svg' else viewer.figure.get_png_data
+            logger.debug("queueing screenshot")
+            data = await queue_screenshot_async(viewer.figure, get_image_data_method)
+            logger.debug("got data, saving to file {filename}")
+            save_to_file(data)
+            logger.debug("saved to file {filename}")
         else:
             raise ValueError(f"Unsupported filetype={filetype} for save_figure")