Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 27 additions & 0 deletions src/wokwi_client/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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:
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What's the use case for this API?

IMHO if there's no good use case, let's not expose it, since it is can be confusing when you have both framebuffer_read and framebuffer_png_bytes.

"""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:
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Let's keep the naming convention consistent. If other methods are called "save_framebuffer_...", then this one should be "get_framebuffer_png_bytes" (or read_...)

"""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:
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

IMHO that doesn't belong to the API - it's better if users implement their own comparison logic as they want (and byte-by-byte comparison may break when we upgrade our PNG library, for example).

"""Compare the current framebuffer with a reference PNG file."""
return await compare_framebuffer_png(
self._transport, id=id, reference=reference, save_mismatch=save_mismatch
)
93 changes: 93 additions & 0 deletions src/wokwi_client/framebuffer.py
Original file line number Diff line number Diff line change
@@ -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
Loading