fix: exporting a large figure hangs

src/py/kaleido/_kaleido_tab/_tab.py

@@ -4,6 +4,7 @@
 from typing import TYPE_CHECKING
 
 import logistro
+import orjson
 
 from . import _devtools_utils as _dtools
 from . import _js_logger
@@ -19,10 +20,18 @@
 
 
 _TEXT_FORMATS = ("svg", "json")  # eps
+_CHUNK_SIZE = 10 * 1024 * 1024  # 10 MB
 
 _logger = logistro.getLogger(__name__)
 
 
+def _orjson_default(obj):
+    """Fallback for types orjson can't handle natively (e.g. NumPy string arrays)."""
+    if hasattr(obj, "tolist"):
+        return obj.tolist()
+    raise TypeError(f"Type is not JSON serializable: {type(obj).__name__}")
+
+
 def _subscribe_new(tab: choreo.Tab, event: str) -> asyncio.Future:
     """Create subscription to tab clearing old ones first: helper function."""
     new_future = tab.subscribe_once(event)
@@ -117,22 +126,38 @@ async def _calc_fig(
         render_prof,
         stepper,
     ) -> bytes:
-        # js script
-        kaleido_js_fn = (
-            r"function(spec, ...args)"
-            r"{"
-            r"return kaleido_scopes.plotly(spec, ...args).then(JSON.stringify);"
-            r"}"
-        )
-        render_prof.profile_log.tick("sending javascript")
-        result = await _dtools.exec_js_fn(
-            self.tab,
-            self._current_js_id,
-            kaleido_js_fn,
+        render_prof.profile_log.tick("serializing spec")
+        spec_str = orjson.dumps(
             spec,
-            topojson,
-            stepper,
-        )
+            default=_orjson_default,
+            option=orjson.OPT_SERIALIZE_NUMPY,
+        ).decode()
+        render_prof.profile_log.tick("spec serialized")
+
+        render_prof.profile_log.tick("sending javascript")
+        if len(spec_str) <= _CHUNK_SIZE:
+            kaleido_js_fn = (
+                r"function(specStr, ...args)"
+                r"{"
+                r"return kaleido_scopes"
+                r".plotly(JSON.parse(specStr), ...args)"
+                r".then(JSON.stringify);"
+                r"}"
+            )
+            result = await _dtools.exec_js_fn(
+                self.tab,
+                self._current_js_id,
+                kaleido_js_fn,
+                spec_str,
+                topojson,
+                stepper,
+            )
+        else:
+            result = await self._calc_fig_chunked(
+                spec_str,
+                topojson=topojson,
+                stepper=stepper,
+            )
         _raise_error(result)
         render_prof.profile_log.tick("javascript sent")
 
@@ -154,3 +179,45 @@ async def _calc_fig(
         render_prof.data_out_size = len(res)
         render_prof.js_log = self.js_logger.log
         return res
+
+    async def _calc_fig_chunked(
+        self,
+        spec_str: str,
+        *,
+        topojson: str | None,
+        stepper,
+    ):
+        _raise_error(
+            await _dtools.exec_js_fn(
+                self.tab,
+                self._current_js_id,
+                r"function() { window.__kaleido_chunks = []; }",
+            )
+        )
+
+        for i in range(0, len(spec_str), _CHUNK_SIZE):
+            chunk = spec_str[i : i + _CHUNK_SIZE]
+            _raise_error(
+                await _dtools.exec_js_fn(
+                    self.tab,
+                    self._current_js_id,
+                    r"function(c) { window.__kaleido_chunks.push(c); }",
+                    chunk,
+                )
+            )
+
+        kaleido_js_fn = (
+            r"function(...args)"
+            r"{"
+            r"var spec = JSON.parse(window.__kaleido_chunks.join(''));"
+            r"delete window.__kaleido_chunks;"
+            r"return kaleido_scopes.plotly(spec, ...args).then(JSON.stringify);"
+            r"}"
+        )
+        return await _dtools.exec_js_fn(
+            self.tab,
+            self._current_js_id,
+            kaleido_js_fn,
+            topojson,
+            stepper,
+        )

src/py/tests/test_large_fig.py

@@ -0,0 +1,26 @@
+import numpy as np
+import plotly.graph_objects as go
+import pytest
+
+import kaleido
+
+TOTAL_POINTS = 5_000_000
+
+
+@pytest.mark.parametrize(
+    ("num_traces", "num_points"),
+    [
+        (1, TOTAL_POINTS),
+        (1_000, TOTAL_POINTS / 1_000),
+    ],
+)
+async def test_large_fig(num_traces, num_points):
+    fig = go.Figure()
+    for _ in range(num_traces):
+        fig.add_trace(
+            go.Scatter(
+                x=np.arange(num_points, dtype=float),
+                y=np.arange(num_points, dtype=float),
+            )
+        )
+    assert isinstance(await kaleido.calc_fig(fig), bytes)