From aa97ff653135b221fd5a74f636622a8d35b926fa Mon Sep 17 00:00:00 2001 From: codeskyblue Date: Wed, 17 Apr 2024 10:56:21 +0800 Subject: [PATCH] fix framebuffer raize ValueError when read size not match --- README.md | 6 +++++- adbutils/_device.py | 4 +++- adbutils/screenshot.py | 9 +++++++-- setup.cfg | 4 ++-- 4 files changed, 17 insertions(+), 6 deletions(-) diff --git a/README.md b/README.md index 568ec2c..4e89733 100644 --- a/README.md +++ b/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) diff --git a/adbutils/_device.py b/adbutils/_device.py index e91ceb1..ef9df31 100644 --- a/adbutils/_device.py +++ b/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 diff --git a/adbutils/screenshot.py b/adbutils/screenshot.py index a815e56..3f8a187 100644 --- a/adbutils/screenshot.py +++ b/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,7 +53,7 @@ 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: @@ -58,5 +61,7 @@ def __screencap(self) -> Image.Image: 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)) \ No newline at end of file diff --git a/setup.cfg b/setup.cfg index 2714fb0..cdefd54 100644 --- a/setup.cfg +++ b/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