diff --git a/uiautomator2/__init__.py b/uiautomator2/__init__.py index 29f0a6d..81a5f94 100644 --- a/uiautomator2/__init__.py +++ b/uiautomator2/__init__.py @@ -236,7 +236,7 @@ def window_size(self): w, h = self._dev.window_size() return w, h - def screenshot(self, filename: Optional[str] = None, format="pillow"): + def screenshot(self, filename: Optional[str] = None, format="pillow", display_id: Optional[int] = None): """ Take screenshot of device @@ -246,13 +246,14 @@ def screenshot(self, filename: Optional[str] = None, format="pillow"): Args: filename (str): saved filename, if filename is set then return None format (str): used when filename is empty. one of ["pillow", "opencv", "raw"] + display_id (int): use specific display if device has multiple screen Examples: screenshot("saved.jpg") screenshot().save("saved.png") cv2.imwrite('saved.jpg', screenshot(format='opencv')) """ - pil_img = self._dev.screenshot() + pil_img = self._dev.screenshot(display_id=display_id) if filename: pil_img.save(filename) return diff --git a/uiautomator2/_selector.py b/uiautomator2/_selector.py index bccb2ac..3d75b2a 100644 --- a/uiautomator2/_selector.py +++ b/uiautomator2/_selector.py @@ -1,6 +1,7 @@ import logging import time import warnings +from typing import Optional import requests from PIL import Image @@ -128,8 +129,8 @@ def info(self): '''ui object info.''' return self.jsonrpc.objInfo(self.selector) - def screenshot(self) -> Image.Image: - im = self.session.screenshot() + def screenshot(self, display_id: Optional[int] = None) -> Image.Image: + im = self.session.screenshot(display_id=display_id) return im.crop(self.bounds()) def click(self, timeout=None, offset=None):