feat: add framebuffer utilities for reading, saving, and comparing PNG images

Author: JakubAndrysek

src/wokwi_client/client.py

@@ -6,6 +6,13 @@
 from pathlib import Path
 from typing import Any, Optional
 
+from wokwi_client.framebuffer import (
+    compare_framebuffer_png,
+    framebuffer_png_bytes,
+    framebuffer_read,
+    save_framebuffer_png,
+)
+
 from .__version__ import get_version
 from .constants import DEFAULT_WS_URL
 from .control import set_control
@@ -251,3 +258,23 @@ async def set_control(
             value: Control value to set (float).
         """
         return await set_control(self._transport, part=part, control=control, value=value)
+
+    async def framebuffer_read(self, id: str) -> ResponseMessage:
+        """Read the current framebuffer for the given device id."""
+        return await framebuffer_read(self._transport, id=id)
+
+    async def framebuffer_png_bytes(self, id: str) -> bytes:
+        """Return the current framebuffer as PNG bytes."""
+        return await framebuffer_png_bytes(self._transport, id=id)
+
+    async def save_framebuffer_png(self, id: str, path: Path, overwrite: bool = True) -> Path:
+        """Save the current framebuffer as a PNG file."""
+        return await save_framebuffer_png(self._transport, id=id, path=path, overwrite=overwrite)
+
+    async def compare_framebuffer_png(
+        self, id: str, reference: Path, save_mismatch: Optional[Path] = None
+    ) -> bool:
+        """Compare the current framebuffer with a reference PNG file."""
+        return await compare_framebuffer_png(
+            self._transport, id=id, reference=reference, save_mismatch=save_mismatch
+        )

src/wokwi_client/client_sync.py

@@ -198,6 +198,41 @@ def set_control(self, part: str, control: str, value: int | bool | float) -> t.A
         assert self._client is not None
         return self._run_async(self._client.set_control(part, control, value))
 
+    def framebuffer_read(self, id: str) -> t.Any:
+        """Read the current framebuffer for the given device id."""
+        if not self._connected:
+            raise RuntimeError("Client not connected")
+        assert self._client is not None
+        return self._run_async(self._client.framebuffer_read(id))
+
+    def framebuffer_png_bytes(self, id: str) -> bytes:
+        """Return the current framebuffer as PNG bytes."""
+        if not self._connected:
+            raise RuntimeError("Client not connected")
+        assert self._client is not None
+        return t.cast(bytes, self._run_async(self._client.framebuffer_png_bytes(id)))
+
+    def save_framebuffer_png(self, id: str, path: Path, overwrite: bool = True) -> Path:
+        """Save the current framebuffer as a PNG file."""
+        if not self._connected:
+            raise RuntimeError("Client not connected")
+        assert self._client is not None
+        return t.cast(
+            Path, self._run_async(self._client.save_framebuffer_png(id, path, overwrite=overwrite))
+        )
+
+    def compare_framebuffer_png(
+        self, id: str, reference: Path, save_mismatch: Path | None = None
+    ) -> bool:
+        """Compare the current framebuffer with a reference PNG file."""
+        if not self._connected:
+            raise RuntimeError("Client not connected")
+        assert self._client is not None
+        return t.cast(
+            bool,
+            self._run_async(self._client.compare_framebuffer_png(id, reference, save_mismatch)),
+        )
+
     @property
     def version(self) -> str:
         if self._client:

src/wokwi_client/framebuffer.py

@@ -0,0 +1,93 @@
+"""Framebuffer command helpers.
+
+Provides utilities to interact with devices exposing a framebuffer (e.g. LCD
+modules) via the `framebuffer:read` command.
+
+Exposed helpers:
+* framebuffer_read          -> raw response (contains base64 PNG at result.png)
+* framebuffer_png_bytes     -> decoded PNG bytes
+* save_framebuffer_png      -> save PNG to disk
+* compare_framebuffer_png   -> compare current framebuffer against reference
+"""
+
+# SPDX-FileCopyrightText: 2025-present CodeMagic LTD
+#
+# SPDX-License-Identifier: MIT
+
+from __future__ import annotations
+
+import base64
+from pathlib import Path
+
+from .exceptions import WokwiError
+from .protocol_types import ResponseMessage
+from .transport import Transport
+
+__all__ = [
+    "framebuffer_read",
+    "framebuffer_png_bytes",
+    "save_framebuffer_png",
+    "compare_framebuffer_png",
+]
+
+
+async def framebuffer_read(transport: Transport, *, id: str) -> ResponseMessage:
+    """Issue `framebuffer:read` for the given device id and return raw response."""
+    return await transport.request("framebuffer:read", {"id": id})
+
+
+def _extract_png_b64(resp: ResponseMessage) -> str:
+    result = resp.get("result", {})
+    png_b64 = result.get("png")
+    if not isinstance(png_b64, str):  # pragma: no cover - defensive
+        raise WokwiError("Malformed framebuffer:read response: missing 'png' base64 string")
+    return png_b64
+
+
+async def framebuffer_png_bytes(transport: Transport, *, id: str) -> bytes:
+    """Return decoded PNG bytes for the framebuffer of device `id`."""
+    resp = await framebuffer_read(transport, id=id)
+    return base64.b64decode(_extract_png_b64(resp))
+
+
+async def save_framebuffer_png(
+    transport: Transport, *, id: str, path: Path, overwrite: bool = True
+) -> Path:
+    """Save the framebuffer PNG to `path` and return the path.
+
+    Args:
+            transport: Active transport.
+            id: Device id (e.g. "lcd1").
+            path: Destination file path.
+            overwrite: Overwrite existing file (default True). If False and file
+                    exists, raises WokwiError.
+    """
+    if path.exists() and not overwrite:
+        raise WokwiError(f"File already exists and overwrite=False: {path}")
+    data = await framebuffer_png_bytes(transport, id=id)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with open(path, "wb") as f:
+        f.write(data)
+    return path
+
+
+async def compare_framebuffer_png(
+    transport: Transport, *, id: str, reference: Path, save_mismatch: Path | None = None
+) -> bool:
+    """Compare the current framebuffer PNG with a reference file.
+
+    Performs a byte-for-byte comparison. If different and `save_mismatch` is
+    provided, writes the current framebuffer PNG there.
+
+    Returns True if identical, False otherwise.
+    """
+    if not reference.exists():
+        raise WokwiError(f"Reference image does not exist: {reference}")
+    current = await framebuffer_png_bytes(transport, id=id)
+    ref_bytes = reference.read_bytes()
+    if current == ref_bytes:
+        return True
+    if save_mismatch:
+        save_mismatch.parent.mkdir(parents=True, exist_ok=True)
+        save_mismatch.write_bytes(current)
+    return False