fastplotlib · kushalkolar · Jul 6, 2023 · Jul 6, 2023 · Jul 6, 2023 · Jul 6, 2023
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -97,4 +97,6 @@ jobs:
       if: ${{ failure() }}
       with:
         name: screenshot-diffs
-        path: examples/desktop/diffs
+        path: |
+          examples/desktop/diffs
+          examples/notebooks/diffs
diff --git a/.github/workflows/screenshots.yml b/.github/workflows/screenshots.yml
@@ -44,8 +44,11 @@ jobs:
         run: |
           # regenerate screenshots
           REGENERATE_SCREENSHOTS=1 pytest -v examples
+          REGENERATE_SCREENSHOTS=1 pytest --nbmake examples/notebooks/
       - uses: actions/upload-artifact@v3
         if: always()
         with:
           name: screenshots
-          path: examples/desktop/screenshots/
+          path: |
+            examples/desktop/screenshots/
+            examples/notebooks/screenshots/
diff --git a/examples/notebooks/lines_cmap.ipynb b/examples/notebooks/lines_cmap.ipynb
@@ -13,6 +13,19 @@
     "import fastplotlib as fpl"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5d2ef4aa-0e4c-4694-ae2e-05da1153a413",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "# this is only for testing, you do not need this to use fastplotlib\n",
+    "from nb_test_utils import plot_test, notebook_finished"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -49,6 +62,18 @@
     "plot.show()"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "727282c3-aadf-420f-a88e-9dd4d4e91263",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "plot_test(\"lines-cmap-white\", plot)"
+   ]
+  },
   {
    "cell_type": "markdown",
    "id": "889b1858-ed64-4d6b-96ad-3883fbe4d38e",
@@ -69,6 +94,19 @@
     "plot.graphics[0].cmap = \"jet\""
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3c9b0bc8-b176-425c-8036-63dc55ab7466",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "# for testing, ignore\n",
+    "plot_test(\"lines-cmap-jet\", plot)"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -81,6 +119,19 @@
     "plot.graphics[0].cmap.values = sine[:, 1]"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6b19d2d4-90e7-40ed-afb9-13abe5474ace",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "# for testing, ignore\n",
+    "plot_test(\"lines-cmap-jet-values\", plot)"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -93,6 +144,19 @@
     "plot.graphics[0].cmap.values = cosine[:, 1]"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "0a6c4739-fa61-4532-865e-21107eab76f9",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "# for testing, ignore\n",
+    "plot_test(\"lines-cmap-jet-values-cosine\", plot)"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -105,6 +169,19 @@
     "plot.graphics[0].cmap = \"viridis\""
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "45acfd2f-09f5-418c-bca5-3e574348b7d5",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "# for testing, ignore\n",
+    "plot_test(\"lines-cmap-viridis\", plot)"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -129,6 +206,19 @@
     "plot.graphics[0].cmap.values = cmap_values"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "7548407f-05ed-4c47-93cc-131c61f8e242",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "# for testing, ignore\n",
+    "plot_test(\"lines-cmap-viridis-values\", plot)"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -147,6 +237,30 @@
    "id": "c290c642-ba5f-4a46-9a17-c434cb39de26",
    "metadata": {},
    "outputs": [],
+   "source": [
+    "# for testing, ignore\n",
+    "plot_test(\"lines-cmap-tab-10\", plot)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c4b9e735-72e9-4f0e-aa3e-43db57e65c99",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "# for testing, ignore\n",
+    "notebook_finished()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f6735cc0-910c-4854-ac50-8ee553a6475e",
+   "metadata": {},
+   "outputs": [],
    "source": []
   }
  ],

diff --git a/examples/notebooks/nb_test_utils.py b/examples/notebooks/nb_test_utils.py
@@ -0,0 +1,87 @@
+from typing import *
+import os
+from pathlib import Path
+
+import imageio.v3 as iio
+import numpy as np
+
+from fastplotlib import Plot, GridPlot
+
+# make dirs for screenshots and diffs
+current_dir = Path(__file__).parent
+
+SCREENSHOTS_DIR = current_dir.joinpath("screenshots")
+DIFFS_DIR = current_dir.joinpath("diffs")
+
+os.makedirs(SCREENSHOTS_DIR, exist_ok=True)
+os.makedirs(DIFFS_DIR, exist_ok=True)
+
+
+# store all the failures to allow the nb to proceed to test other examples
+FAILURES = list()
+
+
+def plot_test(name, plot: Union[Plot, GridPlot]):
+    snapshot = plot.canvas.snapshot()
+
+    if "REGENERATE_SCREENSHOTS" in os.environ.keys():
+        if os.environ["REGENERATE_SCREENSHOTS"] == "1":
+            regenerate_screenshot(name, snapshot.data)
+
+    try:
+        assert_screenshot_equal(name, snapshot.data)
+    except AssertionError:
+        FAILURES.append(name)
+
+
+def regenerate_screenshot(name, data):
+        iio.imwrite(SCREENSHOTS_DIR.joinpath(f"nb-{name}.png"), data)
+
+
+def assert_screenshot_equal(name, data):
+    ground_truth = iio.imread(SCREENSHOTS_DIR.joinpath(f"nb-{name}.png"))
+
+    is_similar = np.allclose(data, ground_truth)
+
+    update_diffs(name, is_similar, data, ground_truth)
+
+    assert is_similar, (
+        f"notebook snapshot for {name} has changed"
+    )
+
+
+def update_diffs(name, is_similar, img, ground_truth):
+    diffs_rgba = None
+
+    def get_diffs_rgba(slicer):
+        # lazily get and cache the diff computation
+        nonlocal diffs_rgba
+        if diffs_rgba is None:
+            # cast to float32 to avoid overflow
+            # compute absolute per-pixel difference
+            diffs_rgba = np.abs(ground_truth.astype("f4") - img)
+            # magnify small values, making it easier to spot small errors
+            diffs_rgba = ((diffs_rgba / 255) ** 0.25) * 255
+            # cast back to uint8
+            diffs_rgba = diffs_rgba.astype("u1")
+        return diffs_rgba[..., slicer]
+
+    # split into an rgb and an alpha diff
+    diffs = {
+        DIFFS_DIR.joinpath(f"nb-diff-{name}-rgb.png"): slice(0, 3),
+        DIFFS_DIR.joinpath(f"nb-diff-{name}-alpha.png"): 3,
+    }
+
+    for path, slicer in diffs.items():
+        if not is_similar:
+            diff = get_diffs_rgba(slicer)
+            iio.imwrite(path, diff)
+        elif path.exists():
+            path.unlink()
+
+
+def notebook_finished():
+    if len(FAILURES) > 0:
+        raise AssertionError(
+            f"Failures for plots:\n{FAILURES}"
+        )
diff --git a/examples/notebooks/scatter.ipynb b/examples/notebooks/scatter.ipynb
diff --git a/examples/notebooks/screenshots/nb-astronaut.png b/examples/notebooks/screenshots/nb-astronaut.png
diff --git a/examples/notebooks/screenshots/nb-astronaut_RGB.png b/examples/notebooks/screenshots/nb-astronaut_RGB.png
diff --git a/examples/notebooks/screenshots/nb-camera.png b/examples/notebooks/screenshots/nb-camera.png
diff --git a/examples/notebooks/screenshots/nb-lines-3d.png b/examples/notebooks/screenshots/nb-lines-3d.png
diff --git a/examples/notebooks/screenshots/nb-lines-cmap-jet-values-cosine.png b/examples/notebooks/screenshots/nb-lines-cmap-jet-values-cosine.png
diff --git a/examples/notebooks/screenshots/nb-lines-cmap-jet-values.png b/examples/notebooks/screenshots/nb-lines-cmap-jet-values.png
diff --git a/examples/notebooks/screenshots/nb-lines-cmap-jet.png b/examples/notebooks/screenshots/nb-lines-cmap-jet.png
diff --git a/examples/notebooks/screenshots/nb-lines-cmap-tab-10.png b/examples/notebooks/screenshots/nb-lines-cmap-tab-10.png
diff --git a/examples/notebooks/screenshots/nb-lines-cmap-viridis-values.png b/examples/notebooks/screenshots/nb-lines-cmap-viridis-values.png
diff --git a/examples/notebooks/screenshots/nb-lines-cmap-viridis.png b/examples/notebooks/screenshots/nb-lines-cmap-viridis.png
diff --git a/examples/notebooks/screenshots/nb-lines-cmap-white.png b/examples/notebooks/screenshots/nb-lines-cmap-white.png
diff --git a/examples/notebooks/screenshots/nb-lines-colors.png b/examples/notebooks/screenshots/nb-lines-colors.png
diff --git a/examples/notebooks/screenshots/nb-lines-data.png b/examples/notebooks/screenshots/nb-lines-data.png
diff --git a/examples/notebooks/screenshots/nb-lines-underlay.png b/examples/notebooks/screenshots/nb-lines-underlay.png
diff --git a/examples/notebooks/screenshots/nb-lines.png b/examples/notebooks/screenshots/nb-lines.png