NF webgl views animation gui

.github/workflows/run_tests.yml

@@ -65,7 +65,7 @@ jobs:
       # A backstop only: pytest-timeout (pytest.ini) already fails a single
       # hung test after 240 s. The suite itself takes 25-30 min on a hosted
       # runner, mostly headless webgl renders.
-      timeout-minutes: 35
+      timeout-minutes: 40
       run: pytest --cov=./
 
     - name: Upload coverage to Codecov

.gitignore

@@ -81,3 +81,4 @@ examples/quickstart/S1_retinotopy.hdf
 # Claude Code working directory (per-worktree scratch: launchers, verify
 # scripts, plans). Not part of the project.
 .claude
+uv.lock

cortex/export/headless.py

@@ -41,9 +41,15 @@
 import concurrent.futures
 import contextlib
 import logging
+import os
+import queue
+import shutil
+import tempfile
 import threading
 import time
-from typing import Any, Mapping, Optional
+from typing import Any, Callable, Mapping, Optional, TypeVar
+
+T = TypeVar("T")
 
 import cortex
 from .. import dataset
@@ -147,16 +153,30 @@ class _PlaywrightThread:
 
     Usage::
 
-        pw_thread = _PlaywrightThread()
+        owns_download_dir = download_dir is None
+    if download_dir is None:
+        download_dir = tempfile.mkdtemp(prefix="pycortex-downloads-")
+    pw_thread = _PlaywrightThread(download_dir=download_dir)
         pw_thread.start(url, timeout=60)   # blocks until page is loaded
         # ... use the pycortex handle (which talks via Tornado, not Playwright) ...
         pw_thread.shutdown()               # tears down browser + playwright
     """
 
-    def __init__(self) -> None:
+    def __init__(self, download_dir: Optional[str] = None) -> None:
         self._ready_future: concurrent.futures.Future[None] = (
             concurrent.futures.Future()
         )
+        # Downloads the page starts -- the animation panel's rendered movies,
+        # the "Save image" button. The listener only queues them and the poll
+        # loop saves them, both on the worker thread: Download.save_as blocks
+        # until the file is complete, which the sync API cannot do from inside
+        # an event callback. Without a directory they are left unsaved.
+        self._download_dir = download_dir
+        self._pending_downloads: list[Any] = []
+        self._downloads: list[str] = []
+        self._downloads_changed = threading.Condition()
+        # Functions queued by run_on_page, run by the poll loop on the worker.
+        self._page_calls: "queue.Queue[tuple[Callable[[Any], Any], concurrent.futures.Future[Any]]]" = queue.Queue()
         self._shutdown_event = threading.Event()
         self._error: Optional[BaseException] = None
         self._thread: Optional[threading.Thread] = None
@@ -213,6 +233,54 @@ def browser_errors(self) -> list[str]:
         with self._errors_lock:
             return list(self._browser_errors)
 
+    @property
+    def downloads(self) -> list[str]:
+        """Paths of the downloads saved so far, oldest first.  Thread-safe."""
+        with self._downloads_changed:
+            return list(self._downloads)
+
+    def wait_for_download(self, timeout: float = 60.0, count: int = 1) -> str:
+        """Wait until `count` downloads have been saved; return the last one's path.
+
+        Raises
+        ------
+        TimeoutError
+            If fewer than `count` downloads arrive within `timeout` seconds.
+        """
+        deadline = time.monotonic() + timeout
+        with self._downloads_changed:
+            while len(self._downloads) < count:
+                remaining = deadline - time.monotonic()
+                if remaining <= 0:
+                    raise TimeoutError(
+                        f"Only {len(self._downloads)} of {count} download(s) "
+                        f"arrived within {timeout:.0f} s"
+                    )
+                self._downloads_changed.wait(remaining)
+            return self._downloads[count - 1]
+
+    def run_on_page(self, fn: Callable[[Any], T], timeout: float = 60.0) -> T:
+        """Run ``fn(page)`` on the worker thread and return what it returns.
+
+        Playwright's sync objects belong to the thread that made them, so the
+        page can only be touched from the worker: this hands it `fn` and waits.
+        It is for what the websocket interface cannot do -- clicking an element,
+        dragging with the mouse, taking a screenshot of part of the page.
+
+        Raises
+        ------
+        TimeoutError
+            If `fn` has not finished within `timeout` seconds.
+        """
+        future: "concurrent.futures.Future[T]" = concurrent.futures.Future()
+        self._page_calls.put((fn, future))
+        try:
+            return future.result(timeout=timeout)
+        except concurrent.futures.TimeoutError:
+            raise TimeoutError(
+                f"The page did not finish the call within {timeout:.0f} s"
+            ) from None
+
     def shutdown(self) -> None:
         """Signal the worker to tear down Playwright and wait for it to finish."""
         self._shutdown_event.set()
@@ -245,6 +313,7 @@ def _worker(self) -> None:
             # errors that fire during page load (e.g. WebGL failures).
             self._page.on("pageerror", self._on_pageerror)
             self._page.on("console", self._on_console)
+            self._page.on("download", self._on_download)
 
             self._page.goto(
                 self._url,
@@ -268,6 +337,8 @@ def _worker(self) -> None:
         # console messages undelivered until _cleanup(). Poll instead: the cheap
         # round-trip is what makes Playwright dispatch them.
         while not self._shutdown_event.wait(EVENT_POLL_INTERVAL):
+            self._save_downloads()
+            self._run_page_calls()
             try:
                 self._page.evaluate("0")
             except Exception:
@@ -292,6 +363,42 @@ def _on_pageerror(self, error: Any) -> None:
         with self._errors_lock:
             self._browser_errors.append(f"[pageerror] {error}")
 
+    def _on_download(self, download: Any) -> None:
+        """Listener for downloads the page starts; the poll loop saves them."""
+        self._pending_downloads.append(download)
+
+    def _save_downloads(self) -> None:
+        """Save queued downloads into the download directory (worker thread)."""
+        while self._pending_downloads:
+            download = self._pending_downloads.pop(0)
+            if self._download_dir is None:
+                continue
+            path = os.path.join(self._download_dir,
+                                os.path.basename(download.suggested_filename))
+            try:
+                download.save_as(path)
+            except Exception:
+                logger.warning("Could not save the download %r",
+                               download.suggested_filename, exc_info=True)
+                continue
+            with self._downloads_changed:
+                self._downloads.append(path)
+                self._downloads_changed.notify_all()
+
+    def _run_page_calls(self) -> None:
+        """Run the functions queued by run_on_page (worker thread)."""
+        while True:
+            try:
+                fn, future = self._page_calls.get_nowait()
+            except queue.Empty:
+                return
+            if not future.set_running_or_notify_cancel():
+                continue
+            try:
+                future.set_result(fn(self._page))
+            except BaseException as exc:  # noqa: BLE001 - handed to the caller
+                future.set_exception(exc)
+
     def _on_console(self, msg: Any) -> None:
         """Listener for console.error / console.warning messages."""
         if msg.type in ("error", "warning"):
@@ -329,6 +436,7 @@ def headless_viewer(
     viewer_params: Mapping[str, Any],
     *,
     timeout: float = 60.0,
+    download_dir: Optional[str] = None,
 ):
     """Context manager that yields a connected ``JSMixer`` handle rendered in a
     headless Chromium browser.
@@ -344,6 +452,12 @@ def headless_viewer(
     timeout : float
         Seconds to wait for the browser to establish the WebSocket connection
         and for ``server.get_client()`` to return (default: 60).
+    download_dir : str or None
+        Directory that files the page downloads -- a movie rendered by the
+        animation panel, an image from "Save image" -- are saved into, as
+        ``handle._pw_thread.downloads``; ``handle._pw_thread.wait_for_download()``
+        waits for one. Default None, meaning a temporary directory removed when
+        the viewer closes.
 
     Yields
     ------
@@ -391,7 +505,10 @@ def headless_viewer(
     #    does not require a GPU or display server, making it usable in
     #    CI / Docker / notebooks.
     # ------------------------------------------------------------------
-    pw_thread = _PlaywrightThread()
+    owns_download_dir = download_dir is None
+    if download_dir is None:
+        download_dir = tempfile.mkdtemp(prefix="pycortex-downloads-")
+    pw_thread = _PlaywrightThread(download_dir=download_dir)
 
     handle = None
     # ------------------------------------------------------------------
@@ -495,3 +612,6 @@ def _await_client() -> None:
             server.stop()
         except Exception:
             logger.warning("Failed to stop Tornado server", exc_info=True)
+
+        if owns_download_dir:
+            shutil.rmtree(download_dir, ignore_errors=True)