From d262d5540e85d51eb4ae5fb11f414164cea6ad74 Mon Sep 17 00:00:00 2001
From: ykh <614457662@qq.com>
Date: Mon, 27 May 2024 17:24:56 +0800
Subject: [PATCH] feat:allow set display_id for screenshot api (#982)

Co-authored-by: yinkanghong <yinkanghong@xiaomi.com>
---
 uiautomator2/__init__.py  | 5 +++--
 uiautomator2/_selector.py | 5 +++--
 2 files changed, 6 insertions(+), 4 deletions(-)

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):