diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index 810abd8..c621100 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -1,4 +1,5 @@ from unittest import TestCase, mock +import io from vncdotool import client, rfb @@ -70,12 +71,24 @@ def test_captureScreen(self, Deferred): d.addCallback.assert_called_once_with(cli._captureSave, fname) assert cli.framebufferUpdateRequest.called + @mock.patch("vncdotool.client.Deferred") + def test_captureScreen_with_format(self, Deferred): + cli = self.client + cli._packet = bytearray(self.MSG_HANDSHAKE) + cli._handleInitial() + cli._handleServerInit(self.MSG_INIT) + cli.vncConnectionMade() + buffer = io.BytesIO() + d = cli.captureScreen(buffer, format="png") + d.addCallback.assert_called_once_with(cli._captureSave, buffer, format="png") + assert cli.framebufferUpdateRequest.called + def test_captureSave(self) -> None: cli = self.client cli.screen = mock.Mock() fname = 'foo.png' r = cli._captureSave(cli.screen, fname) - cli.screen.save.assert_called_once_with(fname) + cli.screen.save.assert_called_once_with(fname, format=None) assert r == cli @mock.patch('PIL.Image.open') diff --git a/vncdotool/client.py b/vncdotool/client.py index 129cfbc..7c7a7cd 100644 --- a/vncdotool/client.py +++ b/vncdotool/client.py @@ -253,10 +253,24 @@ def mouseUp(self: TClient, button: int) -> TClient: return self - def captureScreen(self, fp: TFile, incremental: bool = False) -> Deferred: - """Save the current display to filename""" + def captureScreen( + self, fp: TFile, incremental: bool = False, format: str | None = None + ) -> Deferred: + """ + Capture and save the current VNC screen display to a file. + + Parameters: + fp (TFile): The destination where the screenshot will be saved. + It can be a string path, a `pathlib.Path` object, or a file-like object opened in binary mode. + incremental (bool, optional): + - `False` (default): Captures the entire screen. + - `True`: Captures only the regions of the screen that have changed since the last capture. + format (str | None, optional): + - See Pillow's list of image formats: https://pillow.readthedocs.io/en/stable/handbook/image-file-formats.html + - If set to `None`, Pillow will determine the format based on the provided file name. + """ log.debug("captureScreen %s", fp) - return self._capture(fp, incremental) + return self._capture(fp, incremental, format=format) def captureRegion( self, fp: TFile, x: int, y: int, w: int, h: int, incremental: bool = False @@ -270,19 +284,24 @@ def refreshScreen(self, incremental: bool = False) -> Deferred: self.framebufferUpdateRequest(incremental=incremental) return d - def _capture(self, fp: TFile, incremental: bool, *args: int) -> Deferred: + def _capture( + self, fp: TFile, incremental: bool, *args: int, format: str | None = None + ) -> Deferred: d = self.refreshScreen(incremental) - d.addCallback(self._captureSave, fp, *args) + kwargs = {"format": format} if format else {} + d.addCallback(self._captureSave, fp, *args, **kwargs) return d - def _captureSave(self: TClient, data: object, fp: TFile, *args: int) -> TClient: + def _captureSave( + self: TClient, data: object, fp: TFile, *args: int, format: str | None = None + ) -> TClient: log.debug("captureSave %s", fp) assert self.screen is not None if args: capture = self.screen.crop(args) # type: ignore[arg-type] else: capture = self.screen - capture.save(fp) + capture.save(fp, format=format) return self