From f35c0aafc197858e2b53b09365c02a123092fefd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Antonio=20V=C3=A1zquez=20Blanco?= Date: Wed, 28 Oct 2026 09:34:51 +0100 Subject: [PATCH] Add get_frame and get_frame_bytes for single screenshots The MJPEG streams wait for frames indefinitely, so a screenshot taken without a video signal never returns. get_frame_bytes() returns one raw JPEG frame and get_frame() the decoded image; both close the stream once the frame arrives and raise TimeoutError when none arrives within `timeout` seconds (10 by default). A stream that ends before any frame raises NanoKVMInvalidResponseError. --- README.md | 5 +- docs/api.md | 8 +- docs/sync.md | 4 +- nanokvm/components/stream.py | 36 ++++++++ tests/test_stream.py | 173 ++++++++++++++++++++++++++++------- 5 files changed, 185 insertions(+), 41 deletions(-) diff --git a/README.md b/README.md index ca445bf..0184ca5 100644 --- a/README.md +++ b/README.md @@ -24,9 +24,8 @@ async def main() -> None: await client.paste_text("Hello\nworld!") await client.mouse_click(MouseButton.LEFT, 0.5, 0.5) - async for frame in client.mjpeg_stream(): - frame.save("screen.jpg") - break + frame = await client.get_frame() + frame.save("screen.jpg") asyncio.run(main()) diff --git a/docs/api.md b/docs/api.md index a50a64b..7e2bf79 100644 --- a/docs/api.md +++ b/docs/api.md @@ -117,7 +117,7 @@ while loop.time() < end: | Methods | Requires | |---|---| -| `mjpeg_stream`, `mjpeg_stream_bytes`, `open_ws` | all | +| `mjpeg_stream`, `mjpeg_stream_bytes`, `get_frame`, `get_frame_bytes`, `open_ws` | all | | `set_screen`, `set_frame_detect`, `pause_frame_detect` | non-Pro | | `get_input_region`, `set_input_region`, `get_input_resolution` | non-Pro 2.5.1 | | `reset_hdmi` | PCIe 2.1.5 | @@ -131,7 +131,11 @@ while loop.time() < end: | `set_fps` | Pro 1.2.8 | `mjpeg_stream()` yields decoded `PIL.Image` frames; `mjpeg_stream_bytes()` -yields the raw JPEG bytes. `open_ws(path)` opens a new WebSocket on a path +yields the raw JPEG bytes. `get_frame()` and `get_frame_bytes()` return a +single frame, for a screenshot, and close the stream. They raise +`TimeoutError` when no frame arrives within `timeout` seconds (10 by default), +for example without a video signal; the streams wait for frames indefinitely. +`open_ws(path)` opens a new WebSocket on a path relative to the API base, such as `stream/h264/webrtc`, for protocols the client does not wrap. The caller closes it before the client exits. diff --git a/docs/sync.md b/docs/sync.md index 79bd083..76d5232 100644 --- a/docs/sync.md +++ b/docs/sync.md @@ -15,9 +15,7 @@ with NanoKVMSyncClient("https://kvm.local/api/") as kvm: kvm.push_button(GpioType.POWER, duration_ms=6000) print(kvm.get_gpio().pwr) - for frame in kvm.mjpeg_stream_bytes(): - Path("screen.jpg").write_bytes(frame) - break + Path("screen.jpg").write_bytes(kvm.get_frame_bytes()) ``` - Every coroutine method of `NanoKVMClient` is a blocking method with the same diff --git a/nanokvm/components/stream.py b/nanokvm/components/stream.py index a0cd09c..e02b6d8 100644 --- a/nanokvm/components/stream.py +++ b/nanokvm/components/stream.py @@ -13,10 +13,13 @@ from aiohttp import BodyPartReader, ClientWebSocketResponse, MultipartReader, hdrs from PIL import Image +from ..exceptions import NanoKVMInvalidResponseError from .session import Controller _LOGGER = logging.getLogger(__name__) +_FRAME_TIMEOUT_SECONDS = 10.0 + class StreamController(Controller): """Own MJPEG frame transport and JPEG decoding.""" @@ -62,6 +65,39 @@ async def mjpeg_stream(self) -> AsyncGenerator[Image.Image, None]: session.check_session_generation(generation) yield image + async def get_frame_bytes( + self, *, timeout: float = _FRAME_TIMEOUT_SECONDS + ) -> bytes: + """Get one raw JPEG frame from the MJPEG stream, undecoded. + + ``timeout`` bounds the whole call, connecting included, so a device + without a video signal raises ``TimeoutError`` instead of waiting + forever. The stream is closed once the frame has arrived. + """ + try: + async with asyncio.timeout(timeout) as deadline: + async with contextlib.aclosing(self.mjpeg_stream_bytes()) as frames: + async for data in frames: + return data + except TimeoutError: + if not deadline.expired(): + raise + raise TimeoutError( + f"No video frame arrived within {timeout:g} seconds" + ) from None + raise NanoKVMInvalidResponseError("The MJPEG stream ended without a frame") + + async def get_frame( + self, *, timeout: float = _FRAME_TIMEOUT_SECONDS + ) -> Image.Image: + """Get one decoded frame from the MJPEG stream. + + ``timeout`` bounds the wait for the frame, as in ``get_frame_bytes``. + """ + data = await self.get_frame_bytes(timeout=timeout) + loop = asyncio.get_running_loop() + return await loop.run_in_executor(None, self._parse_jpeg_from_bytes, data) + async def open_ws(self, path: str, **kwargs: Any) -> ClientWebSocketResponse: """Open a new, unshared WebSocket on ``path``, relative to the API base. diff --git a/tests/test_stream.py b/tests/test_stream.py index 2807d3c..8310edb 100644 --- a/tests/test_stream.py +++ b/tests/test_stream.py @@ -1,16 +1,54 @@ """Tests for MJPEG streaming.""" import asyncio +from collections.abc import AsyncGenerator, Awaitable, Callable +from contextlib import asynccontextmanager import io +import time from aiohttp import web from aioresponses import aioresponses from PIL import Image +import pytest -from nanokvm.client import NanoKVMClient +from nanokvm.client import NanoKVMClient, NanoKVMInvalidResponseError from .common import BASE_URL +_MJPEG_HEADERS = {"Content-Type": "multipart/x-mixed-replace; boundary=frame"} + + +def _jpeg(size: tuple[int, int] = (2, 2)) -> bytes: + image_buffer = io.BytesIO() + Image.new("RGB", size).save(image_buffer, format="JPEG") + return image_buffer.getvalue() + + +def _part(data: bytes) -> bytes: + return ( + b"--frame\r\n" + b"Content-Type: image/jpeg\r\n" + + f"Content-Length: {len(data)}\r\n\r\n".encode() + + data + + b"\r\n" + ) + + +@asynccontextmanager +async def _mjpeg_server( + handler: Callable[[web.Request], Awaitable[web.StreamResponse]], +) -> AsyncGenerator[str, None]: + application = web.Application() + application.router.add_get("/api/stream/mjpeg", handler) + runner = web.AppRunner(application) + await runner.setup() + try: + site = web.TCPSite(runner, "127.0.0.1", 0) + await site.start() + yield f"http://127.0.0.1:{runner.addresses[0][1]}/api/" + finally: + await runner.cleanup() + def test_parse_jpeg_returns_fully_loaded_image() -> None: """JPEG parsing returns an image detached from the source stream.""" @@ -26,50 +64,28 @@ def test_parse_jpeg_returns_fully_loaded_image() -> None: async def test_mjpeg_stream_allows_frames_after_request_timeout() -> None: """An MJPEG stream remains open beyond the normal request timeout.""" - image_buffer = io.BytesIO() - Image.new("RGB", (2, 2)).save(image_buffer, format="JPEG") - image_data = image_buffer.getvalue() + image_data = _jpeg() async def stream_handler( request: web.Request, ) -> web.StreamResponse: - response = web.StreamResponse( - headers={"Content-Type": "multipart/x-mixed-replace; boundary=frame"} - ) + response = web.StreamResponse(headers=_MJPEG_HEADERS) await response.prepare(request) for _ in range(2): - await response.write( - b"--frame\r\n" - b"Content-Type: image/jpeg\r\n" - + f"Content-Length: {len(image_data)}\r\n\r\n".encode() - + image_data - + b"\r\n" - ) + await response.write(_part(image_data)) await asyncio.sleep(0.55) await response.write(b"--frame--\r\n") await response.write_eof() return response - application = web.Application() - application.router.add_get("/api/stream/mjpeg", stream_handler) - runner = web.AppRunner(application) - await runner.setup() - site = web.TCPSite(runner, "127.0.0.1", 0) - await site.start() - server_sockets = getattr(site._server, "sockets", None) - assert server_sockets - port = server_sockets[0].getsockname()[1] - frames = [] - try: - async with NanoKVMClient( - f"http://127.0.0.1:{port}/api/", token="test-token", request_timeout=1 - ) as client: - async with asyncio.timeout(3): - async for frame in client.mjpeg_stream(): - frames.append(frame) - finally: - await runner.cleanup() + async with ( + _mjpeg_server(stream_handler) as url, + NanoKVMClient(url, token="test-token", request_timeout=1) as client, + ): + async with asyncio.timeout(3): + async for frame in client.mjpeg_stream(): + frames.append(frame) assert len(frames) == 2 @@ -86,3 +102,94 @@ async def test_mjpeg_stream_bytes_yields_raw_frames(mocked: aioresponses) -> Non frames = [frame async for frame in c.mjpeg_stream_bytes()] assert frames == [b"\xff\xd8one", b"\xff\xd8two"] + + +async def test_get_frame_bytes_returns_first_frame_and_closes_stream() -> None: + """One frame is returned and the device stops streaming right after it.""" + disconnected = asyncio.Event() + + async def endless_handler(request: web.Request) -> web.StreamResponse: + response = web.StreamResponse(headers=_MJPEG_HEADERS) + await response.prepare(request) + try: + for index in range(1000): + await response.write(_part(b"\xff\xd8frame-%d" % index)) + await asyncio.sleep(0.01) + except (ConnectionError, asyncio.CancelledError): + disconnected.set() + raise + return response + + async with ( + _mjpeg_server(endless_handler) as url, + NanoKVMClient(url, token="test-token") as client, + ): + frame = await client.get_frame_bytes() + async with asyncio.timeout(2): + await disconnected.wait() + + assert frame == b"\xff\xd8frame-0" + + +async def test_get_frame_decodes_the_frame(mocked: aioresponses) -> None: + """The decoded variant returns a loaded image.""" + mocked.get( + f"{BASE_URL}stream/mjpeg", + headers=_MJPEG_HEADERS, + body=_part(_jpeg((4, 3))) + b"--frame--\r\n", + ) + async with NanoKVMClient(BASE_URL, token="test-token") as client: + image = await client.get_frame() + + assert image.size == (4, 3) + assert getattr(image, "fp", None) is None + + +@pytest.mark.parametrize("phase", ["headers", "first frame"]) +async def test_get_frame_bytes_times_out_without_a_frame(phase: str) -> None: + """No frame within the timeout raises instead of waiting forever.""" + release = asyncio.Event() + + async def silent_handler(request: web.Request) -> web.StreamResponse: + response = web.StreamResponse(headers=_MJPEG_HEADERS) + if phase == "first frame": + await response.prepare(request) + await release.wait() + return response + + async with ( + _mjpeg_server(silent_handler) as url, + NanoKVMClient(url, token="test-token", request_timeout=30) as client, + ): + started = time.monotonic() + with pytest.raises(TimeoutError, match="within 0.2 seconds"): + await client.get_frame_bytes(timeout=0.2) + elapsed = time.monotonic() - started + release.set() + + assert elapsed < 2 + + +async def test_get_frame_bytes_keeps_other_timeouts( + client: NanoKVMClient, monkeypatch: pytest.MonkeyPatch +) -> None: + """A transport timeout before the deadline is not reported as no frame.""" + + async def failing_stream() -> AsyncGenerator[bytes, None]: + raise TimeoutError("connect timed out") + yield b"" + + monkeypatch.setattr(client, "mjpeg_stream_bytes", failing_stream) + + with pytest.raises(TimeoutError, match="connect timed out"): + await client.get_frame_bytes(timeout=5) + + +async def test_get_frame_bytes_rejects_a_stream_without_frames( + mocked: aioresponses, +) -> None: + """A stream that ends before any frame is an invalid response.""" + mocked.get(f"{BASE_URL}stream/mjpeg", headers=_MJPEG_HEADERS, body=b"--frame--\r\n") + async with NanoKVMClient(BASE_URL, token="test-token") as client: + with pytest.raises(NanoKVMInvalidResponseError, match="without a frame"): + await client.get_frame_bytes()