Skip to content

Commit

Permalink
fix framebuffer raize ValueError when read size not match
Browse files Browse the repository at this point in the history
  • Loading branch information
codeskyblue committed Apr 17, 2024
1 parent fe3aa9c commit aa97ff6
Show file tree
Hide file tree
Showing 4 changed files with 17 additions and 6 deletions.
6 changes: 5 additions & 1 deletion README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
4 changes: 3 additions & 1 deletion adbutils/_device.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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

Expand Down
9 changes: 7 additions & 2 deletions adbutils/screenshot.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -17,6 +18,7 @@
# fix for py37
UnidentifiedImageError = OSError

logger = logging.getLogger(__name__)

class AbstractDevice(abc.ABC):
@property
Expand All @@ -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()
Expand All @@ -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))
4 changes: 2 additions & 2 deletions setup.cfg
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down

0 comments on commit aa97ff6

Please sign in to comment.