From c099bbc025c7bb576292a646b30f14f233fdc35c Mon Sep 17 00:00:00 2001 From: codeskyblue Date: Fri, 9 Aug 2024 19:17:46 +0800 Subject: [PATCH] add fallback for screenshot (#1018) --- uiautomator2/__init__.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/uiautomator2/__init__.py b/uiautomator2/__init__.py index 11e5ebf..36f9671 100644 --- a/uiautomator2/__init__.py +++ b/uiautomator2/__init__.py @@ -252,8 +252,12 @@ def screenshot(self, filename: Optional[str] = None, format="pillow", display_id """ if display_id is None: base64_data = self.jsonrpc.takeScreenshot(1, 80) - jpg_raw = base64.b64decode(base64_data) - pil_img = Image.open(io.BytesIO(jpg_raw)) + # takeScreenshot may return None + if base64_data: + jpg_raw = base64.b64decode(base64_data) + pil_img = Image.open(io.BytesIO(jpg_raw)) + else: + pil_img = self._dev.screenshot(display_id=display_id) else: pil_img = self._dev.screenshot(display_id=display_id)