feat:allow set display_id for screenshot api (#982)

Co-authored-by: yinkanghong <yinkanghong@xiaomi.com>

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

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