From 2d6143d24455586c5a5fbca28e1bdea4bcb593a5 Mon Sep 17 00:00:00 2001
From: codeskyblue <codeskyblue@gmail.com>
Date: Sun, 28 Apr 2024 17:58:38 +0800
Subject: [PATCH] support multi display, disable framebuffer function

---
 adbutils/screenshot.py | 36 ++++++++++++++++++++++++------------
 adbutils/shell.py      |  8 ++++++--
 2 files changed, 30 insertions(+), 14 deletions(-)

diff --git a/adbutils/screenshot.py b/adbutils/screenshot.py
index 3f8a187..336f513 100644
--- a/adbutils/screenshot.py
+++ b/adbutils/screenshot.py
@@ -42,12 +42,16 @@ class ScreenshotExtesion(AbstractDevice):
     def __init__(self):
         self.__framebuffer_ok = True
 
-    def screenshot(self) -> Image.Image:
+    def screenshot(self, display_id: Optional[int] = None) -> Image.Image:
         """ Take a screenshot and return PIL.Image.Image object
-        If capture failed, return a black image
+        Args:
+            display_id: int, default None, see "dumpsys SurfaceFlinger --display-id" for valid display IDs
+        
+        Returns:
+            PIL.Image.Image object, If capture failed, return a black image
         """
         try:
-            pil_image = self.__screencap()
+            pil_image = self.__screencap(display_id)
             if pil_image.mode == "RGBA":
                 pil_image = pil_image.convert("RGB")
             return pil_image
@@ -55,13 +59,21 @@ def screenshot(self) -> Image.Image:
             wsize = self.window_size()
             return Image.new("RGB", wsize, (0, 0, 0))
     
-    def __screencap(self) -> Image.Image:
-        if self.__framebuffer_ok:
-            try:
-                return self.framebuffer()
-            except NotImplementedError:
-                self.__framebuffer_ok = False
-            except UnidentifiedImageError as e:
-                logger.warning("framebuffer error: %s", e)
-        png_bytes = self.shell('screencap -p', encoding=None)
+    def __screencap(self, display_id: int = None) -> Image.Image:
+        """ Take a screenshot and return PIL.Image.Image object
+        """
+        # framebuffer is not stable, so we disable it
+        # MemoryError may occur when using framebuffer
+        
+        # if self.__framebuffer_ok and display_id is None:
+        #     try:
+        #         return self.framebuffer()
+        #     except NotImplementedError:
+        #         self.__framebuffer_ok = False
+        #     except UnidentifiedImageError as e:
+        #         logger.warning("framebuffer error: %s", e)
+        cmdargs = ['screencap', '-p']
+        if display_id is not None:
+            cmdargs.extend(['-d', str(display_id)])
+        png_bytes = self.shell(cmdargs, encoding=None)
         return Image.open(io.BytesIO(png_bytes))
\ No newline at end of file
diff --git a/adbutils/shell.py b/adbutils/shell.py
index fad8da7..4d0f709 100644
--- a/adbutils/shell.py
+++ b/adbutils/shell.py
@@ -152,19 +152,23 @@ def swipe(self, sx, sy, ex, ey, duration: float = 1.0) -> None:
         x1, y1, x2, y2 = map(str, [sx, sy, ex, ey])
         self.shell(["input", "swipe", x1, y1, x2, y2, str(int(duration * 1000))])
 
-    def click(self, x, y) -> None:
+    def click(self, x, y, display_id: Optional[int] = None) -> None:
         """
         simulate android tap
 
         Args:
             x, y: int
+            display_id: int, default None, see "dumpsys SurfaceFlinger --display-id" for valid display IDs
         """
         if any(map(is_percent, [x, y])):
             w, h = self.window_size()
             x = int(x * w) if is_percent(x) else x
             y = int(y * h) if is_percent(y) else y
         x, y = map(str, [x, y])
-        self.shell(["input", "tap", x, y])
+        cmdargs = ["input"]
+        if display_id is not None:
+            cmdargs.extend(['-d', str(display_id)])
+        self.shell(cmdargs + ["tap", x, y])
 
     def send_keys(self, text: str):
         """