Skip to content

Commit c099bbc

Browse files
authored
add fallback for screenshot (#1018)
1 parent 2267dbf commit c099bbc

File tree

1 file changed

+6
-2
lines changed

1 file changed

+6
-2
lines changed

uiautomator2/__init__.py

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -252,8 +252,12 @@ def screenshot(self, filename: Optional[str] = None, format="pillow", display_id
252252
"""
253253
if display_id is None:
254254
base64_data = self.jsonrpc.takeScreenshot(1, 80)
255-
jpg_raw = base64.b64decode(base64_data)
256-
pil_img = Image.open(io.BytesIO(jpg_raw))
255+
# takeScreenshot may return None
256+
if base64_data:
257+
jpg_raw = base64.b64decode(base64_data)
258+
pil_img = Image.open(io.BytesIO(jpg_raw))
259+
else:
260+
pil_img = self._dev.screenshot(display_id=display_id)
257261
else:
258262
pil_img = self._dev.screenshot(display_id=display_id)
259263

0 commit comments

Comments
 (0)