fix framebuffer raize ValueError when read size not match

Author: codeskyblue

README.md

@@ -200,7 +200,11 @@ d.get_state() # same as adb get-state
 
 Take screenshot
 
-```bash
+```python
+# Method 1 (Recommend)
+pil_image = d.screenshot()
+
+# Method 2
 # adb exec-out screencap -p p.png
 png_data = d.shell("screencap -p", encoding=None)
 pathlib.Path("p.png").write_bytes(png_data)

adbutils/_device.py

@@ -13,7 +13,7 @@
 import typing
 from typing import Optional, Union
 
-from PIL import Image
+from PIL import Image, UnidentifiedImageError
 
 from adbutils._deprecated import DeprecatedExtension
 from adbutils.install import InstallExtension
@@ -337,6 +337,8 @@ def framebuffer(self) -> Image.Image:
         if color_format != 'RGBA' and color_format != 'RGB':
             raise NotImplementedError("Unsupported color format")
         buffer = c.read(size)
+        if len(buffer) != size:
+            raise UnidentifiedImageError("framebuffer size not match", size, len(buffer))
         image = Image.frombytes(color_format, (width, height), buffer)
         return image
 

adbutils/screenshot.py

@@ -6,6 +6,7 @@
 
 import abc
 import io
+import logging
 from typing import Optional, Union
 from adbutils.sync import Sync
 from adbutils._proto import WindowSize
@@ -17,6 +18,7 @@
     # fix for py37
     UnidentifiedImageError = OSError
 
+logger = logging.getLogger(__name__)
 
 class AbstractDevice(abc.ABC):
     @property
@@ -42,6 +44,7 @@ def __init__(self):
 
     def screenshot(self) -> Image.Image:
         """ Take a screenshot and return PIL.Image.Image object
+        If capture failed, return a black image
         """
         try:
             pil_image = self.__screencap()
@@ -50,13 +53,15 @@ def screenshot(self) -> Image.Image:
             return pil_image
         except UnidentifiedImageError as e:
             wsize = self.window_size()
-            return Image.new("RGB", wsize, (220, 120, 100))
+            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
-        png_bytes = self.shell('screencap', encoding=None)
+            except UnidentifiedImageError as e:
+                logger.warning("framebuffer error: %s", e)
+        png_bytes = self.shell('screencap -p', encoding=None)
         return Image.open(io.BytesIO(png_bytes))

setup.cfg

@@ -15,9 +15,9 @@ classifier =
     Intended Audience :: Developers
     Operating System :: POSIX
     Programming Language :: Python :: 3
-    Programming Language :: Python :: 3.6
-    Programming Language :: Python :: 3.7
     Programming Language :: Python :: 3.8
+    Programming Language :: Python :: 3.9
+    Programming Language :: Python :: 3.10
     Topic :: Software Development :: Libraries :: Python Modules
     Topic :: Software Development :: Testing