Add get_frame and get_frame_bytes for single screenshots

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())

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.
 

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

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.
 

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()