test: add unit tests for image_array_to_data_uri

tests/test_plotly_utils/test_data_utils.py

@@ -0,0 +1,88 @@
+"""
+Tests for _plotly_utils.data_utils.image_array_to_data_uri.
+
+Uses the pypng backend so the tests do not require Pillow. The generated
+data URIs are decoded with the vendored pypng Reader so the tests assert on
+the encoded pixels rather than only on the data URI prefix.
+"""
+
+import base64
+
+import numpy as np
+import pytest
+
+from _plotly_utils.data_utils import image_array_to_data_uri
+from _plotly_utils.png import Reader
+
+PNG_PREFIX = "data:image/png;base64,"
+
+
+def decode_data_uri(uri):
+    """Decode a PNG data URI back into an array plus the PNG header info."""
+    assert uri.startswith(PNG_PREFIX)
+    png_bytes = base64.b64decode(uri[len(PNG_PREFIX) :])
+    width, height, rows, info = Reader(bytes=png_bytes).read_flat()
+    if info["greyscale"]:
+        channels = 1
+    elif info["alpha"]:
+        channels = 4
+    else:
+        channels = 3
+    decoded = np.array(list(rows), dtype=np.uint8).reshape(height, width, channels)
+    if channels == 1:
+        decoded = decoded[:, :, 0]
+    return decoded, info
+
+
+def test_greyscale_array_round_trips():
+    img = np.array([[0, 255], [128, 64]], dtype=np.uint8)
+
+    decoded, info = decode_data_uri(image_array_to_data_uri(img, backend="pypng"))
+
+    assert info["greyscale"] is True
+    assert info["alpha"] is False
+    np.testing.assert_array_equal(decoded, img)
+
+
+def test_rgb_array_round_trips():
+    img = np.array(
+        [[[255, 0, 0], [0, 255, 0]], [[0, 0, 255], [10, 20, 30]]], dtype=np.uint8
+    )
+
+    decoded, info = decode_data_uri(image_array_to_data_uri(img, backend="pypng"))
+
+    assert info["greyscale"] is False
+    assert info["alpha"] is False
+    np.testing.assert_array_equal(decoded, img)
+
+
+def test_rgba_array_round_trips():
+    img = np.array(
+        [[[255, 0, 0, 255], [0, 255, 0, 128]], [[0, 0, 255, 0], [10, 20, 30, 40]]],
+        dtype=np.uint8,
+    )
+
+    decoded, info = decode_data_uri(image_array_to_data_uri(img, backend="pypng"))
+
+    assert info["greyscale"] is False
+    assert info["alpha"] is True
+    np.testing.assert_array_equal(decoded, img)
+
+
+@pytest.mark.parametrize("compression", [-1, 10])
+def test_invalid_compression_raises(compression):
+    img = np.zeros((2, 2), dtype=np.uint8)
+    with pytest.raises(ValueError, match="compression level"):
+        image_array_to_data_uri(img, backend="pypng", compression=compression)
+
+
+def test_invalid_shape_raises():
+    img = np.zeros(5, dtype=np.uint8)
+    with pytest.raises(ValueError, match="Invalid image shape"):
+        image_array_to_data_uri(img, backend="pypng")
+
+
+def test_jpg_without_pil_backend_raises():
+    img = np.zeros((2, 2), dtype=np.uint8)
+    with pytest.raises(ValueError, match="jpg binary strings"):
+        image_array_to_data_uri(img, backend="pypng", ext="jpg")