From a14634c38907e59ff3ddb3de6deb66a863f13650 Mon Sep 17 00:00:00 2001 From: matejcik Date: Tue, 25 Aug 2020 13:14:47 +0200 Subject: [PATCH] python: make PIL optional in the toif module, add size check --- python/src/trezorlib/debuglink.py | 14 -------------- python/src/trezorlib/toif.py | 28 +++++++++++++++++----------- 2 files changed, 17 insertions(+), 25 deletions(-) diff --git a/python/src/trezorlib/debuglink.py b/python/src/trezorlib/debuglink.py index 320fa300a..cd8222fac 100644 --- a/python/src/trezorlib/debuglink.py +++ b/python/src/trezorlib/debuglink.py @@ -446,20 +446,6 @@ class TrezorClientDebugLink(TrezorClient): def _raw_read(self): __tracebackhide__ = True # for pytest # pylint: disable=W0612 - # if SCREENSHOT and self.debug: - # from PIL import Image - - # layout = self.debug.state().layout - # im = Image.new("RGB", (128, 64)) - # pix = im.load() - # for x in range(128): - # for y in range(64): - # rx, ry = 127 - x, 63 - y - # if (ord(layout[rx + (ry / 8) * 128]) & (1 << (ry % 8))) > 0: - # pix[x, y] = (255, 255, 255) - # im.save("scr%05d.png" % self.screenshot_id) - # self.screenshot_id += 1 - resp = super()._raw_read() resp = self._filter_message(resp) self._check_request(resp) diff --git a/python/src/trezorlib/toif.py b/python/src/trezorlib/toif.py index 10658d747..82b990ab0 100644 --- a/python/src/trezorlib/toif.py +++ b/python/src/trezorlib/toif.py @@ -4,12 +4,13 @@ from typing import Sequence, Tuple import attr +from . import firmware + try: from PIL import Image except ImportError: Image = None -from . import firmware RGBPixel = Tuple[int, int, int] @@ -66,19 +67,14 @@ class Toif: size = attr.ib() # type: Tuple[int, int] data = attr.ib() # type: bytes - def _expected_data_length(self) -> int: + @data.validator + def check_data_size(self, _, value): width, height = self.size if self.mode is firmware.ToifMode.grayscale: - return width * height // 2 + expected_size = width * height // 2 else: - return width * height * 2 - - def to_image(self) -> "Image": - if Image is None: - raise RuntimeError("PIL is not available. Please install via 'pip install Pillow'") - + expected_size = width * height * 2 uncompressed = _decompress(self.data) - expected_size = self._expected_data_length() if len(uncompressed) != expected_size: raise ValueError( "Uncompressed data is {} bytes, expected {}".format( @@ -86,6 +82,14 @@ class Toif: ) ) + def to_image(self) -> "Image": + if Image is None: + raise RuntimeError( + "PIL is not available. Please install via 'pip install Pillow'" + ) + + uncompressed = _decompress(self.data) + if self.mode is firmware.ToifMode.grayscale: pil_mode = "L" raw_data = _to_grayscale(uncompressed) @@ -118,7 +122,9 @@ def load(filename: str) -> Toif: def from_image(image: "Image", background=(0, 0, 0, 255)) -> Toif: if Image is None: - raise RuntimeError("PIL is not available. Please install via 'pip install Pillow'") + raise RuntimeError( + "PIL is not available. Please install via 'pip install Pillow'" + ) if image.mode == "RGBA": background = Image.new("RGBA", image.size, background)