Raise NanoKVMInvalidResponseError for empty snapshots, use get_snapshot in examples

README.md

@@ -11,6 +11,7 @@ pip install nanokvm
 
 ```python
 import asyncio
+from pathlib import Path
 
 from nanokvm.client import NanoKVMClient
 from nanokvm.models import GpioType, MouseButton
@@ -24,9 +25,7 @@ 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
+        Path("screen.jpg").write_bytes(await client.get_snapshot())
 
 
 asyncio.run(main())

docs/api.md

@@ -135,7 +135,8 @@ while loop.time() < end:
 | Methods | Requires |
 |---|---|
 | `mjpeg_stream`, `mjpeg_stream_bytes`, `open_ws`, `open_webrtc_signaling` | all |
-| `get_snapshot`, `set_screen`, `set_frame_detect`, `pause_frame_detect` | non-Pro |
+| `get_snapshot` | non-Pro, or Pro with `interrupt_webrtc=True` |
+| `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 |
 | `get_hdmi_state`, `enable_hdmi`, `disable_hdmi` | PCIe 2.2.8 |
@@ -151,7 +152,7 @@ while loop.time() < end:
 raw JPEG bytes, and `get_snapshot()` the first one. `open_webrtc_signaling()`
 returns the H.264 WebRTC signaling session for the detected hardware.
 [Video streaming](streaming.md) covers both, and why the Pro refuses
-snapshots.
+snapshots by default.
 
 `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

docs/streaming.md

@@ -6,7 +6,7 @@ leaves the media to a WebRTC peer, such as a browser or `aiortc`.
 
 | Need | Method |
 |---|---|
-| One still image | `get_snapshot()` (non-Pro) |
+| One still image | `get_snapshot()`, ending WebRTC on the Pro |
 | Raw JPEG frames | `mjpeg_stream_bytes()` |
 | Decoded `PIL.Image` frames | `mjpeg_stream()` |
 | Live H.264 video, and audio on the Pro | `open_webrtc_signaling()` |
@@ -18,9 +18,13 @@ The Pro encodes one stream mode at a time (`StreamMode`, set with
 WebRTC session in progress stops: its iteration raises
 `NanoKVMVideoModeError`. For that reason `get_snapshot()` raises
 `NanoKVMNotSupportedError` on the Pro without opening the stream
-([NanoKVM-Pro#157](https://github.com/sipeed/NanoKVM-Pro/issues/157)).
-`mjpeg_stream()` and `mjpeg_stream_bytes()` still work on the Pro, at that
-cost.
+([NanoKVM-Pro#157](https://github.com/sipeed/NanoKVM-Pro/issues/157)),
+unless the caller accepts that cost with `get_snapshot(interrupt_webrtc=True)`.
+`mjpeg_stream()` and `mjpeg_stream_bytes()` work on the Pro at the same cost.
+
+The interrupted session does not recover, but a new one started after the
+snapshot gets video. If new sessions stay black, `set_stream_mode()` with the
+WebRTC mode in use switches the encoder back.
 
 ## MJPEG
 
@@ -36,7 +40,8 @@ async with aclosing(client.mjpeg_stream_bytes()) as frames:
 ```
 
 `get_snapshot(timeout=20.0)` returns the first frame, raises `TimeoutError`
-when none arrives in time, and `NanoKVMError` when the stream ends without one.
+when none arrives in time, and `NanoKVMInvalidResponseError` when the stream
+ends without one.
 The streams run until the caller stops iterating; `aclosing()` closes the HTTP
 response as soon as the loop is left.
 

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_snapshot())
 ```
 
 - Every coroutine method of `NanoKVMClient` is a blocking method with the same

nanokvm/components/stream.py

@@ -13,8 +13,7 @@
 from aiohttp import BodyPartReader, ClientWebSocketResponse, MultipartReader, hdrs
 from PIL import Image
 
-from ..compatibility import require_hardware
-from ..exceptions import NanoKVMError
+from ..exceptions import NanoKVMInvalidResponseError, NanoKVMNotSupportedError
 from ..models.common import HWFamily
 from ..webrtc import WebRTCSignaling
 from .session import Controller
@@ -55,21 +54,31 @@ async def mjpeg_stream_bytes(self) -> AsyncGenerator[bytes, None]:
                 session.check_session_generation(generation)
                 yield data
 
-    @require_hardware(HWFamily.NON_PRO)
-    async def get_snapshot(self, timeout: float = 20.0) -> bytes:
+    async def get_snapshot(
+        self, timeout: float = 20.0, *, interrupt_webrtc: bool = False
+    ) -> bytes:
         """Return one raw JPEG frame from the MJPEG stream.
 
-        Not on the Pro: its MJPEG and H.264 streams share one encoder mode, so
-        opening MJPEG interrupts any WebRTC session
-        (https://github.com/sipeed/NanoKVM-Pro/issues/157).
+        On the Pro, MJPEG and H.264 share one encoder mode, so opening MJPEG
+        ends any WebRTC session in progress
+        (https://github.com/sipeed/NanoKVM-Pro/issues/157). The Pro refuses
+        the snapshot unless ``interrupt_webrtc`` accepts that.
         """
+        session = self._session
+        if session._hw_version is None:
+            await session.detect_hardware()
+        if session.is_hardware_family(HWFamily.PRO) and not interrupt_webrtc:
+            raise NanoKVMNotSupportedError(
+                "get_snapshot on the Pro ends any WebRTC session in progress; "
+                "pass interrupt_webrtc=True to accept that"
+            )
         async with (
             asyncio.timeout(timeout),
             contextlib.aclosing(self.mjpeg_stream_bytes()) as frames,
         ):
             async for frame in frames:
                 return frame
-        raise NanoKVMError("MJPEG stream ended without a frame")
+        raise NanoKVMInvalidResponseError("MJPEG stream ended without a frame")
 
     async def mjpeg_stream(self) -> AsyncGenerator[Image.Image, None]:
         """Stream decoded MJPEG frames."""

tests/test_stream.py

@@ -9,7 +9,7 @@
 import pytest
 
 from nanokvm.client import NanoKVMClient
-from nanokvm.exceptions import NanoKVMError, NanoKVMNotSupportedError
+from nanokvm.exceptions import NanoKVMInvalidResponseError, NanoKVMNotSupportedError
 from nanokvm.models import HWVersion
 
 from .common import BASE_URL, mark_detected
@@ -112,14 +112,28 @@ async def test_get_snapshot_without_frames_raises(mocked: aioresponses) -> None:
     mocked.get(f"{BASE_URL}stream/mjpeg", headers=_MJPEG_HEADERS, body=b"--frame--\r\n")
     async with NanoKVMClient(BASE_URL, token="test-token") as c:
         mark_detected(c)
-        with pytest.raises(NanoKVMError, match="without a frame"):
+        with pytest.raises(NanoKVMInvalidResponseError, match="without a frame"):
             await c.get_snapshot()
 
 
 async def test_get_snapshot_is_refused_on_the_pro(mocked: aioresponses) -> None:
     """The Pro never opens MJPEG for a snapshot, so WebRTC stays up."""
     async with NanoKVMClient(BASE_URL, token="test-token") as c:
         mark_detected(c, HWVersion.PRO)
-        with pytest.raises(NanoKVMNotSupportedError):
+        with pytest.raises(NanoKVMNotSupportedError, match="interrupt_webrtc"):
             await c.get_snapshot()
     assert not mocked.requests
+
+
+async def test_get_snapshot_on_the_pro_when_interrupting_webrtc(
+    mocked: aioresponses,
+) -> None:
+    """The Pro takes a snapshot once the caller accepts ending WebRTC."""
+    mocked.get(
+        f"{BASE_URL}stream/mjpeg",
+        headers=_MJPEG_HEADERS,
+        body=b"--frame\r\nContent-Type: image/jpeg\r\n\r\n\xff\xd8one\r\n--frame--\r\n",
+    )
+    async with NanoKVMClient(BASE_URL, token="test-token") as c:
+        mark_detected(c, HWVersion.PRO)
+        assert await c.get_snapshot(interrupt_webrtc=True) == b"\xff\xd8one"