add fallback for screenshot (#1018)

openatx · Aug 9, 2024 · c099bbc · c099bbc
1 parent 2267dbf
commit c099bbc
Showing 1 changed file with 6 additions and 2 deletions.
diff --git 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)