diff --git a/.github/workflows/python-package.yml b/.github/workflows/python-package.yml index 5989973c..ae4d71d5 100644 --- a/.github/workflows/python-package.yml +++ b/.github/workflows/python-package.yml @@ -15,7 +15,7 @@ env: jobs: build: - runs-on: ubuntu-latest + runs-on: ubuntu-22.04 strategy: matrix: python-version: ['3.7', '3.8', '3.9', '3.10', '3.11', '3.12'] diff --git a/androidtv/adb_manager/adb_manager_async.py b/androidtv/adb_manager/adb_manager_async.py index e5325291..db84c596 100644 --- a/androidtv/adb_manager/adb_manager_async.py +++ b/androidtv/adb_manager/adb_manager_async.py @@ -439,6 +439,33 @@ async def shell(self, cmd): _LOGGER.debug("Sending command to %s:%d via adb-shell: %s", self.host, self.port, cmd) return await self._adb.shell(cmd) + async def shell_binary(self, cmd): + """Send an ADB command using the Python ADB implementation and return raw bytes. + + Parameters + ---------- + cmd : str + The ADB command to be sent + + Returns + ------- + bytes, None + The raw response from the device, if there is a response + + """ + if not self.available: + _LOGGER.debug( + "ADB binary command not sent to %s:%d because adb-shell connection is not established: %s", + self.host, + self.port, + cmd, + ) + return None + + async with _acquire(self._adb_lock): + _LOGGER.debug("Sending binary command to %s:%d via adb-shell: %s", self.host, self.port, cmd) + return await self._adb.shell(cmd, decode=False) + class ADBServerAsync(object): """A manager for ADB connections that uses an ADB server. diff --git a/androidtv/adb_manager/adb_manager_sync.py b/androidtv/adb_manager/adb_manager_sync.py index e93283d3..b0c1d054 100644 --- a/androidtv/adb_manager/adb_manager_sync.py +++ b/androidtv/adb_manager/adb_manager_sync.py @@ -321,6 +321,33 @@ def shell(self, cmd): _LOGGER.debug("Sending command to %s:%d via adb-shell: %s", self.host, self.port, cmd) return self._adb.shell(cmd) + def shell_binary(self, cmd): + """Send an ADB command using the Python ADB implementation and return raw bytes. + + Parameters + ---------- + cmd : str + The ADB command to be sent + + Returns + ------- + bytes, None + The raw response from the device, if there is a response + + """ + if not self.available: + _LOGGER.debug( + "ADB binary command not sent to %s:%d because adb-shell connection is not established: %s", + self.host, + self.port, + cmd, + ) + return None + + with _acquire(self._adb_lock): + _LOGGER.debug("Sending binary command to %s:%d via adb-shell: %s", self.host, self.port, cmd) + return self._adb.shell(cmd, decode=False) + class ADBServerSync(object): """A manager for ADB connections that uses an ADB server. diff --git a/androidtv/basetv/basetv_async.py b/androidtv/basetv/basetv_async.py index d901d10c..100d741b 100644 --- a/androidtv/basetv/basetv_async.py +++ b/androidtv/basetv/basetv_async.py @@ -3,11 +3,20 @@ ADB Debugging must be enabled. """ +import asyncio import logging +from time import monotonic from .basetv import BaseTV from .. import constants from ..adb_manager.adb_manager_async import ADBPythonAsync, ADBServerAsync +from ..framebuffer import ( + GZIP_SCREENCAP_COMMAND, + RAW_SCREENCAP_COMMAND, + ScreencapError, + decode_framebuffer, + decode_gzip_screencap, +) _LOGGER = logging.getLogger(__name__) @@ -85,6 +94,7 @@ def __init__( adb = ADBServerAsync(host, port, adb_server_ip, adb_server_port) BaseTV.__init__(self, adb, host, port, adbkey, adb_server_ip, adb_server_port, state_detection_rules) + self._screencap_method = None # ======================================================================= # # # @@ -145,16 +155,108 @@ async def adb_push(self, local_path, device_path): async def adb_screencap(self): """Take a screencap. - This calls :py:meth:`androidtv.adb_manager.adb_manager_async.ADBPythonAsync.screencap` or :py:meth:`androidtv.adb_manager.adb_manager_async.ADBServerAsync.screencap`, - depending on whether the Python ADB implementation or an ADB server is used for communicating with the device. + Returns the screen as a binary .png image, or ``None`` if the device did not return anything. + + Note: there are three capture paths. Plain ``screencap`` and ``screencap | gzip -1`` are read + back as raw bytes (:py:meth:`androidtv.adb_manager.adb_manager_async.ADBPythonAsync.shell_binary`) + with the PNG encoded on the host (the host-side decode runs in an executor); or the device encodes + the PNG itself via device-side ``screencap -p`` + (:py:meth:`androidtv.adb_manager.adb_manager_async.ADBPythonAsync.screencap` or + :py:meth:`androidtv.adb_manager.adb_manager_async.ADBServerAsync.screencap`, depending on whether the + Python ADB implementation or an ADB server is used for communicating with the device). The ADB + server path has no binary shell, so it uses only device-side ``screencap -p``. Neither path is + always faster. The first call after :meth:`adb_connect` times each binary-safe candidate back to + back, discards any candidate that fails validation, pins the fastest valid one, and returns its + image; later calls use only the pinned method. If the pinned binary method later fails, it falls + back to device-side ``screencap -p`` and pins that for good. Returns ------- - bytes - The screencap as a binary .png image + bytes, None + The screencap as a binary .png image, or ``None`` if the device did not return anything """ - return await self._adb.screencap() + method = getattr(self, "_screencap_method", None) + if method == "native": + return await self._adb.screencap() + if method == "raw": + return await self._raw_screencap() + if method == "gzip": + return await self._gzip_screencap() + return await self._calibrate_screencap() + + async def _binary_shell(self, cmd): + """Run a shell command and return raw bytes, or ``None`` when the path has no binary shell.""" + shell_binary = getattr(self._adb, "shell_binary", None) + if shell_binary is None: + return None + return await shell_binary(cmd) + + async def _raw_screencap(self): + """Run the raw path, falling back to device-side ``screencap -p`` on any decode failure.""" + data = await self._binary_shell(RAW_SCREENCAP_COMMAND) + if not isinstance(data, bytes) or not data: + self._screencap_method = "native" + return await self._adb.screencap() + try: + return await asyncio.get_running_loop().run_in_executor(None, decode_framebuffer, data) + except ScreencapError as err: + _LOGGER.debug("Falling back to screencap -p: %s", err) + self._screencap_method = "native" + return await self._adb.screencap() + + async def _gzip_screencap(self): + """Run the gzip path, falling back to device-side ``screencap -p`` on any decode failure.""" + data = await self._binary_shell(GZIP_SCREENCAP_COMMAND) + if not isinstance(data, bytes) or not data: + self._screencap_method = "native" + return await self._adb.screencap() + try: + return await asyncio.get_running_loop().run_in_executor(None, decode_gzip_screencap, data) + except ScreencapError as err: + _LOGGER.debug("Falling back to screencap -p: %s", err) + self._screencap_method = "native" + return await self._adb.screencap() + + async def _calibrate_screencap(self): + """Time each binary-safe candidate back to back, pin the fastest valid one, return its image.""" + loop = asyncio.get_running_loop() + results = {} + binary_shell = getattr(self._adb, "shell_binary", None) + if binary_shell is None: + _LOGGER.debug("ADB server has no binary shell; using screencap -p") + self._screencap_method = "native" + return await self._adb.screencap() + for name, cmd, decode in ( + ("raw", RAW_SCREENCAP_COMMAND, decode_framebuffer), + ("gzip", GZIP_SCREENCAP_COMMAND, decode_gzip_screencap), + ): + start = monotonic() + data = await binary_shell(cmd) + if isinstance(data, bytes) and data: + try: + image = await loop.run_in_executor(None, decode, data) + except ScreencapError as err: + _LOGGER.debug("Screencap candidate %r failed validation: %s", name, err) + else: + results[name] = (monotonic() - start, image) + continue + _LOGGER.debug("Screencap candidate %r returned no data", name) + start = monotonic() + native_image = await self._adb.screencap() + native_dt = monotonic() - start + if native_image is not None: + results["native"] = (native_dt, native_image) + _LOGGER.debug( + "Screencap calibration: %s", + ", ".join(f"{name} {dt:.3f}s" for name, (dt, _) in sorted(results.items(), key=lambda kv: kv[1][0])), + ) + if not results: + self._screencap_method = "native" + return None + name = min(results, key=lambda key: results[key][0]) + self._screencap_method = name + return results[name][1] async def adb_connect( self, @@ -179,6 +281,7 @@ async def adb_connect( Whether or not the connection was successfully established and the device is available """ + self._screencap_method = None if isinstance(self._adb, ADBPythonAsync): return await self._adb.connect(log_errors, auth_timeout_s, transport_timeout_s) return await self._adb.connect(log_errors) diff --git a/androidtv/basetv/basetv_sync.py b/androidtv/basetv/basetv_sync.py index 57ad0d84..9d29cbc2 100644 --- a/androidtv/basetv/basetv_sync.py +++ b/androidtv/basetv/basetv_sync.py @@ -4,9 +4,17 @@ """ import logging +from time import monotonic from .basetv import BaseTV from .. import constants +from ..framebuffer import ( + GZIP_SCREENCAP_COMMAND, + RAW_SCREENCAP_COMMAND, + ScreencapError, + decode_framebuffer, + decode_gzip_screencap, +) from ..adb_manager.adb_manager_sync import ADBPythonSync, ADBServerSync _LOGGER = logging.getLogger(__name__) @@ -85,6 +93,7 @@ def __init__( adb = ADBServerSync(host, port, adb_server_ip, adb_server_port) BaseTV.__init__(self, adb, host, port, adbkey, adb_server_ip, adb_server_port, state_detection_rules) + self._screencap_method = None # ======================================================================= # # # @@ -145,16 +154,106 @@ def adb_push(self, local_path, device_path): def adb_screencap(self): """Take a screencap. - This calls :py:meth:`androidtv.adb_manager.adb_manager_sync.ADBPythonSync.screencap` or :py:meth:`androidtv.adb_manager.adb_manager_sync.ADBServerSync.screencap`, - depending on whether the Python ADB implementation or an ADB server is used for communicating with the device. + Returns the screen as a binary .png image, or ``None`` if the device did not return anything. + + Note: there are three capture paths. Plain ``screencap`` and ``screencap | gzip -1`` are read + back as raw bytes (:py:meth:`androidtv.adb_manager.adb_manager_sync.ADBPythonSync.shell_binary`) + with the PNG encoded on the host; or the device encodes the PNG itself via device-side + ``screencap -p`` (:py:meth:`androidtv.adb_manager.adb_manager_sync.ADBPythonSync.screencap` + or :py:meth:`androidtv.adb_manager.adb_manager_sync.ADBServerSync.screencap`, depending on whether the + Python ADB implementation or an ADB server is used for communicating with the device). The ADB + server path has no binary shell, so it uses only device-side ``screencap -p``. Neither path is + always faster. The first call after :meth:`adb_connect` times each binary-safe candidate back to + back, discards any candidate that fails validation, pins the fastest valid one, and returns its + image; later calls use only the pinned method. If the pinned binary method later fails, it falls + back to device-side ``screencap -p`` and pins that for good. Returns ------- - bytes - The screencap as a binary .png image + bytes, None + The screencap as a binary .png image, or ``None`` if the device did not return anything """ - return self._adb.screencap() + method = getattr(self, "_screencap_method", None) + if method == "native": + return self._adb.screencap() + if method == "raw": + return self._raw_screencap() + if method == "gzip": + return self._gzip_screencap() + return self._calibrate_screencap() + + def _binary_shell(self, cmd): + """Run a shell command and return raw bytes, or ``None`` when the path has no binary shell.""" + shell_binary = getattr(self._adb, "shell_binary", None) + if shell_binary is None: + return None + return shell_binary(cmd) + + def _raw_screencap(self): + """Run the raw path, falling back to device-side ``screencap -p`` on any decode failure.""" + data = self._binary_shell(RAW_SCREENCAP_COMMAND) + if not isinstance(data, bytes) or not data: + self._screencap_method = "native" + return self._adb.screencap() + try: + return decode_framebuffer(data) + except ScreencapError as err: + _LOGGER.debug("Falling back to screencap -p: %s", err) + self._screencap_method = "native" + return self._adb.screencap() + + def _gzip_screencap(self): + """Run the gzip path, falling back to device-side ``screencap -p`` on any decode failure.""" + data = self._binary_shell(GZIP_SCREENCAP_COMMAND) + if not isinstance(data, bytes) or not data: + self._screencap_method = "native" + return self._adb.screencap() + try: + return decode_gzip_screencap(data) + except ScreencapError as err: + _LOGGER.debug("Falling back to screencap -p: %s", err) + self._screencap_method = "native" + return self._adb.screencap() + + def _calibrate_screencap(self): + """Time each binary-safe candidate back to back, pin the fastest valid one, return its image.""" + results = {} + binary_shell = getattr(self._adb, "shell_binary", None) + if binary_shell is None: + _LOGGER.debug("ADB server has no binary shell; using screencap -p") + self._screencap_method = "native" + return self._adb.screencap() + for name, cmd, decode in ( + ("raw", RAW_SCREENCAP_COMMAND, decode_framebuffer), + ("gzip", GZIP_SCREENCAP_COMMAND, decode_gzip_screencap), + ): + start = monotonic() + data = binary_shell(cmd) + if isinstance(data, bytes) and data: + try: + image = decode(data) + except ScreencapError as err: + _LOGGER.debug("Screencap candidate %r failed validation: %s", name, err) + else: + results[name] = (monotonic() - start, image) + continue + _LOGGER.debug("Screencap candidate %r returned no data", name) + start = monotonic() + native_image = self._adb.screencap() + native_dt = monotonic() - start + if native_image is not None: + results["native"] = (native_dt, native_image) + _LOGGER.debug( + "Screencap calibration: %s", + ", ".join(f"{name} {dt:.3f}s" for name, (dt, _) in sorted(results.items(), key=lambda kv: kv[1][0])), + ) + if not results: + self._screencap_method = "native" + return None + name = min(results, key=lambda key: results[key][0]) + self._screencap_method = name + return results[name][1] def adb_connect( self, @@ -179,6 +278,7 @@ def adb_connect( Whether or not the connection was successfully established and the device is available """ + self._screencap_method = None if isinstance(self._adb, ADBPythonSync): return self._adb.connect(log_errors, auth_timeout_s, transport_timeout_s) return self._adb.connect(log_errors) diff --git a/androidtv/framebuffer.py b/androidtv/framebuffer.py new file mode 100644 index 00000000..868a4355 --- /dev/null +++ b/androidtv/framebuffer.py @@ -0,0 +1,166 @@ +"""Capture Android framebuffers over ADB and encode them as PNG. + +``screencap -p`` makes the device CPU encode the PNG. Two binary alternatives +avoid that: plain ``screencap`` (raw pixels, decoded on the host) and +``screencap | gzip -1`` (raw pixels gzipped on the device). Neither path is +always faster: on one Hi3751V350 projector the medians were 7.4 s for +``screencap -p``, about 2.5 s for raw ``screencap`` read binary (0.7 s grab +plus transfer of 8.3 MB), and about 4.1 s for the gzip path. Callers time each +path back to back on first use and keep the fastest valid one for that +device. The host parses the Android framebuffer header and encodes the PNG +itself using only the standard library. +""" + +import gzip +import struct +import sys +import zlib +from array import array + +RAW_SCREENCAP_COMMAND = "screencap" + +GZIP_SCREENCAP_COMMAND = "screencap | gzip -1" + +PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n" + +# Android ``PixelFormat`` values we can decode: format -> bytes per pixel. +_BYTES_PER_PIXEL = {1: 4, 2: 4, 3: 3, 4: 2, 5: 4} + +_EXPAND_5 = bytes((c << 3) | (c >> 2) for c in range(32)) +_EXPAND_6 = bytes((c << 2) | (c >> 4) for c in range(64)) + + +class ScreencapError(Exception): + """A compressed screencap could not be transferred or decoded.""" + + +class UnsupportedScreencapFormatError(ScreencapError): + """The framebuffer uses a pixel format this decoder does not handle.""" + + +def decode_gzip_screencap(data): + """Decode the raw gzip output of :data:`GZIP_SCREENCAP_COMMAND`. + + Parameters + ---------- + data : bytes + The raw bytes printed by the device shell + + Returns + ------- + bytes + The screencap as a binary .png image + + Raises + ------ + ScreencapError + If the data is not valid gzip or the framebuffer is truncated + UnsupportedScreencapFormatError + If the framebuffer pixel format is not supported + + """ + try: + raw = gzip.decompress(data) + except (ValueError, OSError, EOFError, zlib.error) as err: + raise ScreencapError(f"Corrupt gzip screencap: {err}") from err + return decode_framebuffer(raw) + + +def decode_framebuffer(raw): + """Convert a raw Android framebuffer dump (``screencap`` without ``-p``) to PNG. + + The dump starts with little-endian ``width``, ``height`` and ``format`` + fields. Android 10+ adds a ``dataspace`` field, so the header is either 12 + or 16 bytes long; the header length is inferred from the total length. + + Parameters + ---------- + raw : bytes + The raw framebuffer dump + + Returns + ------- + bytes + The screencap as a binary .png image + + """ + if len(raw) < 12: + raise ScreencapError("Truncated screencap header") + width, height, pixel_format = struct.unpack_from("IIBBBBB", width, height, 8, color_type, 0, 0, 0) + return b"".join( + ( + PNG_SIGNATURE, + _png_chunk(b"IHDR", ihdr), + _png_chunk(b"IDAT", zlib.compress(scanlines, 1)), + _png_chunk(b"IEND", b""), + ) + ) + + +def _png_chunk(tag, data): + crc = zlib.crc32(data, zlib.crc32(tag)) & 0xFFFFFFFF + return struct.pack(">I", len(data)) + tag + data + struct.pack(">I", crc) + + +def _bgra_to_rgba(pixels): + out = bytearray(len(pixels)) + out[0::4] = pixels[2::4] + out[1::4] = pixels[1::4] + out[2::4] = pixels[0::4] + out[3::4] = pixels[3::4] + return out + + +def _rgb565_to_rgb(pixels, count): + """Expand little-endian RGB_565 (red in the top five bits) to 8-bit RGB.""" + values = array("H") + values.frombytes(pixels) + if sys.byteorder == "big": + values.byteswap() + out = bytearray(count * 3) + out[0::3] = bytes(_EXPAND_5[v >> 11] for v in values) + out[1::3] = bytes(_EXPAND_6[(v >> 5) & 0x3F] for v in values) + out[2::3] = bytes(_EXPAND_5[v & 0x1F] for v in values) + return out diff --git a/tests/async_patchers.py b/tests/async_patchers.py index 3ac715f2..18551792 100644 --- a/tests/async_patchers.py +++ b/tests/async_patchers.py @@ -1,5 +1,6 @@ """Define patches used for androidtv tests.""" +from inspect import iscoroutinefunction from unittest.mock import patch try: @@ -8,8 +9,13 @@ from unittest.mock import MagicMock class AsyncMock(MagicMock): + """Await async side effects like unittest.mock.AsyncMock on Python 3.8+.""" + async def __call__(self, *args, **kwargs): - return super(AsyncMock, self).__call__(*args, **kwargs) + result = super(AsyncMock, self).__call__(*args, **kwargs) + if iscoroutinefunction(self.side_effect): + return await result + return result KEY_PYTHON = "python" diff --git a/tests/test_adb_manager_sync.py b/tests/test_adb_manager_sync.py index 482df913..89da9b26 100644 --- a/tests/test_adb_manager_sync.py +++ b/tests/test_adb_manager_sync.py @@ -332,3 +332,33 @@ def test_close(self): self.adb.close() self.assertFalse(self.adb.available) + + +class TestADBPythonSyncShellBinary(unittest.TestCase): + """Test the `ADBPythonSync.shell_binary` method.""" + + PATCH_KEY = "python" + + def setUp(self): + """Create an `ADBPythonSync` instance.""" + with patchers.PATCH_ADB_DEVICE_TCP, patchers.patch_connect(True)[self.PATCH_KEY]: + self.adb = ADBPythonSync("HOST", 5555) + + def test_adb_shell_binary_returns_raw_bytes(self): + """Test that ``shell_binary`` returns the undecoded device output when connected.""" + with patchers.patch_connect(True)[self.PATCH_KEY], patchers.patch_shell(b"\x00\x01\x02")[self.PATCH_KEY]: + self.assertTrue(self.adb.connect()) + self.assertEqual(self.adb.shell_binary("screencap"), b"\x00\x01\x02") + + def test_adb_shell_binary_unavailable_returns_none(self): + """Test that ``shell_binary`` returns ``None`` when the adb-shell connection is not established.""" + self.assertFalse(self.adb.available) + self.assertIsNone(self.adb.shell_binary("screencap")) + + def test_adb_shell_binary_lock_not_acquired(self): + """Test that ``shell_binary`` raises when the ADB lock is held by another caller.""" + with patchers.patch_connect(True)[self.PATCH_KEY], patchers.patch_shell(b"\x00")[self.PATCH_KEY]: + self.assertTrue(self.adb.connect()) + with patch.object(self.adb, "_adb_lock", LockedLock()): + with self.assertRaises(LockNotAcquiredException): + self.adb.shell_binary("screencap") diff --git a/tests/test_basetv_async.py b/tests/test_basetv_async.py index 96f60f88..c74daea7 100644 --- a/tests/test_basetv_async.py +++ b/tests/test_basetv_async.py @@ -59,7 +59,10 @@ async def test_adb_push(self): @awaiter async def test_adb_screencap(self): """Test that the ``adb_screencap`` method works correctly.""" - with patch.object(self.btv._adb, "screencap", return_value=PNG_IMAGE, new_callable=async_patchers.AsyncMock): + # adb_shell returning "" means the device cannot compress; adb_screencap then falls back to _adb.screencap + with patch.object(self.btv, "adb_shell", return_value="", new_callable=async_patchers.AsyncMock), patch.object( + self.btv._adb, "screencap", return_value=PNG_IMAGE, new_callable=async_patchers.AsyncMock + ): self.assertEqual(await self.btv.adb_screencap(), PNG_IMAGE) @awaiter diff --git a/tests/test_basetv_sync.py b/tests/test_basetv_sync.py index ce6e7104..3b980d40 100644 --- a/tests/test_basetv_sync.py +++ b/tests/test_basetv_sync.py @@ -230,7 +230,10 @@ def test_adb_push(self): def test_adb_screencap(self): """Test that the ``adb_screencap`` method works correctly.""" - with patch.object(self.btv._adb, "screencap", return_value=PNG_IMAGE): + # adb_shell returning "" means the device cannot compress; adb_screencap then falls back to _adb.screencap + with patch.object(self.btv, "adb_shell", return_value=""), patch.object( + self.btv._adb, "screencap", return_value=PNG_IMAGE + ): self.assertEqual(self.btv.adb_screencap(), PNG_IMAGE) def test_keys(self): diff --git a/tests/test_framebuffer.py b/tests/test_framebuffer.py new file mode 100644 index 00000000..cf8c3ef0 --- /dev/null +++ b/tests/test_framebuffer.py @@ -0,0 +1,448 @@ +import base64 +import gzip +import struct +import unittest +import zlib +from unittest.mock import Mock, patch + +from androidtv import framebuffer +from androidtv.basetv.basetv_sync import BaseTVSync +from androidtv.basetv.basetv_async import BaseTVAsync +from .async_patchers import AsyncMock +from .async_wrapper import awaiter + + +PNG_IMAGE = b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\n\x00\x00\x00\n\x08\x06\x00\x00\x00\x8d2\xcf\xbd\x00\x00\x00\x04sBIT\x08\x08\x08\x08|\x08d\x88\x00\x00\x00\tpHYs\x00\x00\x0fa\x00\x00\x0fa\x01\xa8?\xa7i\x00\x00\x00\x0eIDAT\x18\x95c`\x18\x05\x83\x13\x00\x00\x01\x9a\x00\x01\x16\xca\xd3i\x00\x00\x00\x00IEND\xaeB`\x82" + + +def _header(width, height, pixel_format, dataspace=None): + if dataspace is None: + return struct.pack("I", png, offset) + tag = png[offset + 4 : offset + 8] + data = png[offset + 8 : offset + 8 + length] + (crc,) = struct.unpack_from(">I", png, offset + 8 + length) + assert crc == zlib.crc32(tag + data) & 0xFFFFFFFF + chunks[tag] = data + offset += 12 + length + width, height, depth, color_type, _, _, _ = struct.unpack(">IIBBBBB", chunks[b"IHDR"]) + assert depth == 8 + return width, height, color_type, zlib.decompress(chunks[b"IDAT"]) + + +def _rows(scanlines, height, stride): + """Strip the per-row filter byte and return the pixel bytes of each row.""" + row_len = stride + 1 + rows = [] + for row in range(height): + line = scanlines[row * row_len : (row + 1) * row_len] + assert line[0] == 0 + rows.append(line[1:]) + return rows + + +class TestDecodeFramebuffer(unittest.TestCase): + def test_rgba_header_12(self): + raw = _header(2, 1, 1) + bytes([1, 2, 3, 4, 5, 6, 7, 8]) + width, height, color_type, scanlines = _parse_png(framebuffer.decode_framebuffer(raw)) + self.assertEqual((width, height, color_type), (2, 1, 6)) + self.assertEqual(_rows(scanlines, 1, 8), [bytes([1, 2, 3, 4, 5, 6, 7, 8])]) + + def test_rgba_header_16_with_dataspace(self): + raw = _header(1, 2, 1, dataspace=0) + bytes(range(8)) + width, height, color_type, scanlines = _parse_png(framebuffer.decode_framebuffer(raw)) + self.assertEqual((width, height, color_type), (1, 2, 6)) + self.assertEqual(_rows(scanlines, 2, 4), [bytes(range(4)), bytes(range(4, 8))]) + + def test_bgra_is_swapped_to_rgba(self): + raw = _header(1, 1, 5) + bytes([10, 20, 30, 40]) + _, _, color_type, scanlines = _parse_png(framebuffer.decode_framebuffer(raw)) + self.assertEqual(color_type, 6) + self.assertEqual(_rows(scanlines, 1, 4), [bytes([30, 20, 10, 40])]) + + def test_rgbx_drops_padding_byte(self): + raw = _header(2, 1, 2) + bytes([1, 2, 3, 0, 4, 5, 6, 0]) + _, _, color_type, scanlines = _parse_png(framebuffer.decode_framebuffer(raw)) + self.assertEqual(color_type, 2) + self.assertEqual(_rows(scanlines, 1, 6), [bytes([1, 2, 3, 4, 5, 6])]) + + def test_rgb888(self): + raw = _header(1, 2, 3) + bytes([1, 2, 3, 4, 5, 6]) + _, _, color_type, scanlines = _parse_png(framebuffer.decode_framebuffer(raw)) + self.assertEqual(color_type, 2) + self.assertEqual(_rows(scanlines, 2, 3), [bytes([1, 2, 3]), bytes([4, 5, 6])]) + + def test_rgb565_expands_to_8_bit(self): + white = struct.pack("