From dab255f8a4082fb32b84a229263349bc9d23897e Mon Sep 17 00:00:00 2001 From: Jakub Andrysek Date: Wed, 20 Aug 2025 10:48:34 +0200 Subject: [PATCH 1/2] feat: add framebuffer command helpers for reading and saving PNGs --- src/wokwi_client/client.py | 27 ++++++++++ src/wokwi_client/framebuffer.py | 93 +++++++++++++++++++++++++++++++++ 2 files changed, 120 insertions(+) create mode 100644 src/wokwi_client/framebuffer.py diff --git a/src/wokwi_client/client.py b/src/wokwi_client/client.py index a1feed6..87fc315 100644 --- a/src/wokwi_client/client.py +++ b/src/wokwi_client/client.py @@ -5,6 +5,13 @@ from pathlib import Path from typing import Any, Optional, Union +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 @@ -233,3 +240,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 + ) diff --git a/src/wokwi_client/framebuffer.py b/src/wokwi_client/framebuffer.py new file mode 100644 index 0000000..ef3ae53 --- /dev/null +++ b/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 From e9321cf010361ef48f7ea6abd7cf653e494647b5 Mon Sep 17 00:00:00 2001 From: Jakub Andrysek Date: Wed, 20 Aug 2025 13:19:50 +0200 Subject: [PATCH 2/2] refactor: rename framebuffer helper functions for clarity --- src/wokwi_client/client.py | 20 +++---------------- src/wokwi_client/framebuffer.py | 35 ++++++--------------------------- 2 files changed, 9 insertions(+), 46 deletions(-) diff --git a/src/wokwi_client/client.py b/src/wokwi_client/client.py index 87fc315..25235cf 100644 --- a/src/wokwi_client/client.py +++ b/src/wokwi_client/client.py @@ -6,9 +6,7 @@ from typing import Any, Optional, Union from wokwi_client.framebuffer import ( - compare_framebuffer_png, - framebuffer_png_bytes, - framebuffer_read, + read_framebuffer_png_bytes, save_framebuffer_png, ) @@ -241,22 +239,10 @@ async def set_control( """ 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: + async def read_framebuffer_png_bytes(self, id: str) -> bytes: """Return the current framebuffer as PNG bytes.""" - return await framebuffer_png_bytes(self._transport, id=id) + return await read_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 - ) diff --git a/src/wokwi_client/framebuffer.py b/src/wokwi_client/framebuffer.py index ef3ae53..ca7f05a 100644 --- a/src/wokwi_client/framebuffer.py +++ b/src/wokwi_client/framebuffer.py @@ -24,14 +24,13 @@ from .transport import Transport __all__ = [ - "framebuffer_read", - "framebuffer_png_bytes", + "read_framebuffer", + "read_framebuffer_png_bytes", "save_framebuffer_png", - "compare_framebuffer_png", ] -async def framebuffer_read(transport: Transport, *, id: str) -> ResponseMessage: +async def read_framebuffer(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}) @@ -44,9 +43,9 @@ def _extract_png_b64(resp: ResponseMessage) -> str: return png_b64 -async def framebuffer_png_bytes(transport: Transport, *, id: str) -> bytes: +async def read_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) + resp = await read_framebuffer(transport, id=id) return base64.b64decode(_extract_png_b64(resp)) @@ -64,30 +63,8 @@ async def save_framebuffer_png( """ if path.exists() and not overwrite: raise WokwiError(f"File already exists and overwrite=False: {path}") - data = await framebuffer_png_bytes(transport, id=id) + data = await read_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