Merge branch 'plotly:master' into patch-3

Author: tschm

src/py/CHANGELOG.txt

@@ -1,3 +1,6 @@
+v1.2.0
+- Try to use plotly JSON encoder instead of default
+
 v1.1.0
 - Add testing
 - Fix a variety of type bugs

src/py/integration_tests/dates/__init__.py

@@ -0,0 +1 @@
+"""These are really just scripts, run them as such. Include pickle dev group."""

src/py/integration_tests/dates/figs/scatter_map_w_layers.pkl.zst

@@ -0,0 +1,58 @@
+"""Create pickled fig for use in integration tests."""
+
+import pickle
+from pathlib import Path
+
+import datashader as ds
+import datashader.transfer_functions as tf
+import pandas as pd
+import plotly.express as px
+import zstandard as zstd
+from colorcet import fire
+
+cctx = zstd.ZstdCompressor(level=20)
+
+df = pd.read_csv(
+    "https://raw.githubusercontent.com/plotly/datasets/master/uber-rides-data1.csv",
+)
+dff = (
+    df.query("Lat < 40.82")
+    .query("Lat > 40.70")
+    .query("Lon > -74.02")
+    .query("Lon < -73.91")
+)
+
+
+cvs = ds.Canvas(plot_width=1000, plot_height=1000)
+agg = cvs.points(dff, x="Lon", y="Lat")
+# agg is an xarray object, see http://xarray.pydata.org/en/stable/ for more details
+coords_lat, coords_lon = agg.coords["Lat"].to_numpy(), agg.coords["Lon"].to_numpy()
+# Corners of the image
+coordinates = [
+    [coords_lon[0], coords_lat[0]],
+    [coords_lon[-1], coords_lat[0]],
+    [coords_lon[-1], coords_lat[-1]],
+    [coords_lon[0], coords_lat[-1]],
+]
+
+
+img = tf.shade(agg, cmap=fire)[::-1].to_pil()
+
+
+# Trick to create rapidly a figure with map axes
+fig = px.scatter_map(dff[:1], lat="Lat", lon="Lon", zoom=12)
+# Add the datashader image as a tile map layer image
+fig.update_layout(
+    map_style="carto-darkmatter",
+    map_layers=[{"sourcetype": "image", "source": img, "coordinates": coordinates}],
+)
+
+raw = pickle.dumps(fig, protocol=5)  # >=3.8
+compressed = cctx.compress(raw)
+with Path(f"./figs/{Path(__file__).stem}.pkl.zst").open("wb") as f:
+    f.write(compressed)
+
+print(  # noqa: T201
+    f"{Path(__file__).stem}.pkl: "
+    f"{len(raw) / 1024:.1f} -> {len(compressed) / 1024:.1f} KB",
+)

src/py/integration_tests/dates/scatter_map_w_layers.py

@@ -28,6 +28,15 @@
 
 _logger = logistro.getLogger(__name__)
 
+try:
+    from plotly.utils import PlotlyJSONEncoder  #  noqa: I001
+    from choreographer import channels
+
+    channels.register_custom_encoder(PlotlyJSONEncoder)
+    _logger.debug("Successfully registered PlotlyJSONEncoder.")
+except ImportError as e:
+    _logger.debug(f'Couldn\'t import plotly due to "{e!s}" - skipping.')
+
 # Show a warning if the installed Plotly version
 # is incompatible with this version of Kaleido
 warn_incompatible_plotly()

src/py/kaleido/kaleido.py

@@ -14,7 +14,7 @@ kaleido = ['vendor/**']
 [project]
 name = "kaleido"
 description = "Plotly graph export library"
-license = {file = "LICENSE.md"}
+license = { "file" = "LICENSE.md" }
 readme = "README.md"
 requires-python = ">=3.8"
 dynamic = ["version"]
@@ -26,7 +26,7 @@ maintainers = [
   {name = "Andrew Pikul", email = "ajpikul@gmail.com"},
 ]
 dependencies = [
-  "choreographer>=1.0.10",
+  "choreographer>=1.1.1",
   "logistro>=1.0.8",
   "orjson>=3.10.15",
   "packaging",
@@ -55,6 +55,13 @@ dev = [
   "typing-extensions>=4.12.2",
   "hypothesis>=6.113.0",
 ]
+pickles = [
+    "colorcet>=3.1.0",
+    "datashader>=0.15.2",
+    "pillow>=10.4.0",
+    "plotly[express]>=6.3.0",
+    "zstandard>=0.23.0",
+]
 
 [tool.ruff.lint]
 select = ["ALL"]

src/py/pyproject.toml

@@ -0,0 +1,115 @@
+import numpy as np
+import plotly.graph_objects as go
+import pytest
+from plotly.subplots import make_subplots
+
+import kaleido
+
+pytestmark = pytest.mark.asyncio(loop_scope="function")
+
+rng = np.random.default_rng()  # creates a Generator instance
+
+
+async def test_complex_plotly_encoder():
+    """Test that kaleido can handle complex Plotly figures with numpy arrays."""
+
+    # Create complex numpy arrays
+    x = np.linspace(0, 4 * np.pi, 100)
+    y1 = np.sin(x) * np.exp(-x / 10)
+
+    # Create a 2D array for heatmap
+    z = np.outer(
+        np.sin(np.linspace(0, np.pi, 20)),
+        np.cos(np.linspace(0, np.pi, 30)),
+    )
+
+    # Create subplot figure
+    fig = make_subplots(
+        rows=2,
+        cols=2,
+        subplot_titles=(
+            "Complex Scatter",
+            "Heatmap",
+            "Bar with Color",
+            "3D Surface",
+        ),
+        specs=[
+            [{"type": "scatter"}, {"type": "heatmap"}],
+            [{"type": "bar"}, {"type": "surface"}],
+        ],
+    )
+
+    # Scatter with numpy marker sizes and complex styling
+    fig.add_trace(
+        go.Scatter(
+            x=x,
+            y=y1,
+            mode="lines+markers",
+            line={"color": "blue", "width": 2},
+            marker={
+                "size": [rng.integers(2, 12) for _ in range(len(x))],
+                "color": np.abs(y1),
+                "colorscale": "Viridis",
+                "showscale": True,
+            },
+        ),
+        row=1,
+        col=1,
+    )
+
+    # Heatmap with 2D numpy array
+    fig.add_trace(
+        go.Heatmap(
+            z=z,
+            colorscale="Plasma",
+        ),
+        row=1,
+        col=2,
+    )
+
+    # Bar chart with numpy data and color mapping
+    categories = np.array(["A", "B", "C", "D", "E", "F"])
+    values = rng.normal(50, 20, len(categories))
+
+    fig.add_trace(
+        go.Bar(
+            x=categories,
+            y=values,
+            marker={
+                "color": np.abs(values),
+                "colorscale": "Blues",
+                "line": {"width": 2, "color": "black"},
+            },
+        ),
+        row=2,
+        col=1,
+    )
+
+    # 3D surface with complex numpy operations
+    x_surf = np.linspace(-3, 3, 30)
+    y_surf = np.linspace(-3, 3, 30)
+    X, Y = np.meshgrid(x_surf, y_surf)  # noqa: N806
+    Z = np.sin(np.sqrt(X**2 + Y**2)) * np.exp(-(X**2 + Y**2) / 10)  # noqa: N806
+
+    fig.add_trace(
+        go.Surface(
+            x=X,
+            y=Y,
+            z=Z,
+            colorscale="Cividis",
+        ),
+        row=2,
+        col=2,
+    )
+
+    # Complex layout with numpy-based annotations
+    fig.update_layout(
+        title="Complex Numpy Figure for Encoder Testing",
+        height=800,
+        width=1200,
+    )
+
+    # Render with kaleido
+    img_bytes = await kaleido.calc_fig(fig)
+
+    assert isinstance(img_bytes, bytes)

src/py/tests/test_encoder_one_off.py

@@ -175,32 +175,29 @@ def st_mathjax(dir_path: Path):
 # Test default combinations
 @pytest.mark.order(1)
 async def test_defaults_no_plotly_available():
-    """Test defaults when plotly package is not available."""
-    if not find_spec("plotly"):
-        raise ImportError("Tests must be run with plotly installed to function")
-
-    old_path = sys.path
-    sys.path = sys.path[:1]
-    if find_spec("plotly"):
-        raise RuntimeError(
-            "Plotly cannot be imported during this test, "
-            "as this tests default behavior while trying to import plotly. "
-            "The best solution is to make sure this test always runs first, "
-            "or if you really need to, run it separately and then skip it "
-            "in the main group.",
-        )
-
-    # Test no imports (plotly not available)
-    no_imports = PageGenerator().generate_index()
-    scripts, _encodings = get_scripts_from_html(no_imports)
-
-    # Should have mathjax, plotly default, and kaleido_scopes
-    assert len(scripts) == 3  # noqa: PLR2004
-    assert scripts[0] == DEFAULT_MATHJAX
-    assert scripts[1] == DEFAULT_PLOTLY
-    assert scripts[2].endswith("kaleido_scopes.js")
+    """
+    Test defaults when plotly package is not available.
 
-    sys.path = old_path
+    When we generate_index(), if we don't have plotly in path, we use a CDN.
+    """
+    _old_path = sys.path
+    try:
+        sys.path = []
+        _plotly_mo = sys.modules.pop("plotly", None)
+
+        # Test no imports (plotly not available)
+        no_imports = PageGenerator().generate_index()
+        scripts, _encodings = get_scripts_from_html(no_imports)
+
+        # Should have mathjax, plotly default, and kaleido_scopes
+        assert len(scripts) == 3  # noqa: PLR2004
+        assert scripts[0] == DEFAULT_MATHJAX
+        assert scripts[1] == DEFAULT_PLOTLY
+        assert scripts[2].endswith("kaleido_scopes.js")
+    finally:
+        sys.path = _old_path
+        if _plotly_mo:
+            sys.modules.update({"plotly": _plotly_mo})
 
 
 async def test_defaults_with_plotly_available():