[Refactor] Break QGFImageFile's _save function into smaller pieces (#21124)

lib/python/qmk/painter_qgf.py

@@ -1,9 +1,11 @@
 # Copyright 2021 Nick Brassel (@tzarc)
+# Copyright 2023 Pablo Martinez (@elpekenin) <elpekenin@elpekenin.dev>
 # SPDX-License-Identifier: GPL-2.0-or-later
 
 # Quantum Graphics File "QGF" Image File Format.
 # See https://docs.qmk.fm/#/quantum_painter_qgf for more information.
 
+import functools
 from colorsys import rgb_to_hsv
 from types import FunctionType
 from PIL import Image, ImageFile, ImageChops
@@ -15,6 +17,12 @@ def o24(i):
     return o16(i & 0xFFFF) + o8((i & 0xFF0000) >> 16)
 
 
+# Helper to convert from RGB888 to the QMK "dialect" of HSV888
+def rgb888_to_qmk_hsv888(e):
+    hsv = rgb_to_hsv(e[0] / 255.0, e[1] / 255.0, e[2] / 255.0)
+    return (int(hsv[0] * 255.0), int(hsv[1] * 255.0), int(hsv[2] * 255.0))
+
+
 ########################################################################################################################
 
 
@@ -60,6 +68,14 @@ class QGFGraphicsDescriptor:
             + o16(self.frame_count)  # frame count
         )
 
+    @property
+    def image_size(self):
+        return self.image_width, self.image_height
+
+    @image_size.setter
+    def image_size(self, size):
+        self.image_width, self.image_height = size
+
 
 ########################################################################################################################
 
@@ -180,6 +196,14 @@ class QGFFrameDeltaDescriptorV1:
                  + o16(self.bottom)  # bottom
                  )
 
+    @property
+    def bbox(self):
+        return self.left, self.top, self.right, self.bottom
+
+    @bbox.setter
+    def bbox(self, bbox):
+        self.left, self.top, self.right, self.bottom = bbox
+
 
 ########################################################################################################################
 
@@ -221,26 +245,10 @@ def _accept(prefix):
     return False
 
 
-def _save(im, fp, filename):
+def _for_all_frames(x: FunctionType, /, images):
-    """Helper method used by PIL to write to an output file.
-    """
-    # Work out from the parameters if we need to do anything special
-    encoderinfo = im.encoderinfo.copy()
-    append_images = list(encoderinfo.get("append_images", []))
-    verbose = encoderinfo.get("verbose", False)
-    use_deltas = encoderinfo.get("use_deltas", True)
-    use_rle = encoderinfo.get("use_rle", True)
-
-    # Helper for inline verbose prints
-    def vprint(s):
-        if verbose:
-            print(s)
-
-    # Helper to iterate through all frames in the input image
-    def _for_all_frames(x: FunctionType):
     frame_num = 0
     last_frame = None
-        for frame in [im] + append_images:
+    for frame in images:
         # Get number of of frames in this image
         nfr = getattr(frame, "n_frames", 1)
         for idx in range(nfr):
@@ -251,44 +259,11 @@ def _save(im, fp, filename):
             last_frame = copy
             frame_num += 1
 
-    # Collect all the frame sizes
-    frame_sizes = []
-    _for_all_frames(lambda idx, frame, last_frame: frame_sizes.append(frame.size))
-
-    # Make sure all frames are the same size
-    if len(list(set(frame_sizes))) != 1:
-        raise ValueError("Mismatching sizes on frames")
-
-    # Write out the initial graphics descriptor (and write a dummy value), so that we can come back and fill in the
-    # correct values once we've written all the frames to the output
-    graphics_descriptor_location = fp.tell()
-    graphics_descriptor = QGFGraphicsDescriptor()
-    graphics_descriptor.frame_count = len(frame_sizes)
-    graphics_descriptor.image_width = frame_sizes[0][0]
-    graphics_descriptor.image_height = frame_sizes[0][1]
-    vprint(f'{"Graphics descriptor block":26s} {fp.tell():5d}d / {fp.tell():04X}h')
-    graphics_descriptor.write(fp)
-
-    # Work out the frame offset descriptor location (and write a dummy value), so that we can come back and fill in the
-    # correct offsets once we've written all the frames to the output
-    frame_offset_location = fp.tell()
-    frame_offsets = QGFFrameOffsetDescriptorV1(graphics_descriptor.frame_count)
-    vprint(f'{"Frame offsets block":26s} {fp.tell():5d}d / {fp.tell():04X}h')
-    frame_offsets.write(fp)
-
-    # Helper function to save each frame to the output file
-    def _write_frame(idx, frame, last_frame):
-        # If we replace the frame we're going to output with a delta, we can override it here
-        this_frame = frame
-        location = (0, 0)
-        size = frame.size
-
-        # Work out the format we're going to use
-        format = encoderinfo["qmk_format"]
 
+def _compress_image(frame, last_frame, *, use_rle, use_deltas, format_, **_kwargs):
     # Convert the original frame so we can do comparisons
-        converted = qmk.painter.convert_requested_format(this_frame, format)
+    converted = qmk.painter.convert_requested_format(frame, format_)
-        graphic_data = qmk.painter.convert_image_bytes(converted, format)
+    graphic_data = qmk.painter.convert_image_bytes(converted, format_)
 
     # Convert the raw data to RLE-encoded if requested
     raw_data = graphic_data[1]
@@ -299,6 +274,7 @@ def _save(im, fp, filename):
 
     # Work out if a delta frame is smaller than injecting it directly
     use_delta_this_frame = False
+    bbox = None
     if use_deltas and last_frame is not None:
         # If we want to use deltas, then find the difference
         diff = ImageChops.difference(frame, last_frame)
@@ -310,12 +286,10 @@ def _save(im, fp, filename):
         if bbox:
             # ...create the delta frame by cropping the original.
             delta_frame = frame.crop(bbox)
-                delta_location = (bbox[0], bbox[1])
-                delta_size = (bbox[2] - bbox[0], bbox[3] - bbox[1])
 
             # Convert the delta frame to the requested format
-                delta_converted = qmk.painter.convert_requested_format(delta_frame, format)
+            delta_converted = qmk.painter.convert_requested_format(delta_frame, format_)
-                delta_graphic_data = qmk.painter.convert_image_bytes(delta_converted, format)
+            delta_graphic_data = qmk.painter.convert_image_bytes(delta_converted, format_)
 
             # Work out how large the delta frame is going to be with compression etc.
             delta_raw_data = delta_graphic_data[1]
@@ -329,10 +303,6 @@ def _save(im, fp, filename):
             # sizing constraints.
             if (len(delta_image_data) + QGFFrameDeltaDescriptorV1.length) < len(image_data):
                 # Copy across all the delta equivalents so that the rest of the processing acts on those
-                    this_frame = delta_frame
-                    location = delta_location
-                    size = delta_size
-                    converted = delta_converted
                 graphic_data = delta_graphic_data
                 raw_data = delta_raw_data
                 rle_data = delta_rle_data
@@ -340,27 +310,51 @@ def _save(im, fp, filename):
                 image_data = delta_image_data
                 use_delta_this_frame = True
 
+        # Default to whole image
+        bbox = bbox or [0, 0, *frame.size]
+        # Fix sze (as per #20296), we need to cast first as tuples are inmutable
+        bbox = list(bbox)
+        bbox[2] -= 1
+        bbox[3] -= 1
+
+    return {
+        "bbox": bbox,
+        "graphic_data": graphic_data,
+        "image_data": image_data,
+        "use_delta_this_frame": use_delta_this_frame,
+        "use_raw_this_frame": use_raw_this_frame,
+    }
+
+
+# Helper function to save each frame to the output file
+def _write_frame(idx, frame, last_frame, *, fp, frame_offsets, **kwargs):
+    # Not an argument of the function as it would consume from **kwargs
+    format_ = kwargs["format_"]
+
+    # (potentially) Apply RLE and/or delta, and work out output image's information
+    outputs = _compress_image(frame, last_frame, **kwargs)
+    bbox = outputs["bbox"]
+    graphic_data = outputs["graphic_data"]
+    image_data = outputs["image_data"]
+    use_delta_this_frame = outputs["use_delta_this_frame"]
+    use_raw_this_frame = outputs["use_raw_this_frame"]
+
     # Write out the frame descriptor
     frame_offsets.frame_offsets[idx] = fp.tell()
     vprint(f'{f"Frame {idx:3d} base":26s} {fp.tell():5d}d / {fp.tell():04X}h')
     frame_descriptor = QGFFrameDescriptorV1()
     frame_descriptor.is_delta = use_delta_this_frame
     frame_descriptor.is_transparent = False
-        frame_descriptor.format = format['image_format_byte']
+    frame_descriptor.format = format_['image_format_byte']
     frame_descriptor.compression = 0x00 if use_raw_this_frame else 0x01  # See qp.h, painter_compression_t
-        frame_descriptor.delay = frame.info['duration'] if 'duration' in frame.info else 1000  # If we're not an animation, just pretend we're delaying for 1000ms
+    frame_descriptor.delay = frame.info.get('duration', 1000)  # If we're not an animation, just pretend we're delaying for 1000ms
     frame_descriptor.write(fp)
 
     # Write out the palette if required
-        if format['has_palette']:
+    if format_['has_palette']:
         palette = graphic_data[0]
         palette_descriptor = QGFFramePaletteDescriptorV1()
 
-            # Helper to convert from RGB888 to the QMK "dialect" of HSV888
-            def rgb888_to_qmk_hsv888(e):
-                hsv = rgb_to_hsv(e[0] / 255.0, e[1] / 255.0, e[2] / 255.0)
-                return (int(hsv[0] * 255.0), int(hsv[1] * 255.0), int(hsv[2] * 255.0))
-
         # Convert all palette entries to HSV888 and write to the output
         palette_descriptor.palette_entries = list(map(rgb888_to_qmk_hsv888, palette))
         vprint(f'{f"Frame {idx:3d} palette":26s} {fp.tell():5d}d / {fp.tell():04X}h')
@@ -370,10 +364,7 @@ def _save(im, fp, filename):
     if use_delta_this_frame:
         # Set up the rendering location of where the delta frame should be situated
         delta_descriptor = QGFFrameDeltaDescriptorV1()
-            delta_descriptor.left = location[0]
+        delta_descriptor.bbox = bbox
-            delta_descriptor.top = location[1]
-            delta_descriptor.right = location[0] + size[0] - 1
-            delta_descriptor.bottom = location[1] + size[1] - 1
 
         # Write the delta frame to the output
         vprint(f'{f"Frame {idx:3d} delta":26s} {fp.tell():5d}d / {fp.tell():04X}h')
@@ -385,8 +376,49 @@ def _save(im, fp, filename):
     vprint(f'{f"Frame {idx:3d} data":26s} {fp.tell():5d}d / {fp.tell():04X}h')
     data_descriptor.write(fp)
 
+
+def _save(im, fp, _filename):
+    """Helper method used by PIL to write to an output file.
+    """
+    # Work out from the parameters if we need to do anything special
+    encoderinfo = im.encoderinfo.copy()
+
+    # Helper for prints, noop taking any args if not verbose
+    global vprint
+    verbose = encoderinfo.get("verbose", False)
+    vprint = print if verbose else lambda *_args, **_kwargs: None
+
+    # Helper to iterate through all frames in the input image
+    append_images = list(encoderinfo.get("append_images", []))
+    for_all_frames = functools.partial(_for_all_frames, images=[im, *append_images])
+
+    # Collect all the frame sizes
+    frame_sizes = []
+    for_all_frames(lambda _idx, frame, _last_frame: frame_sizes.append(frame.size))
+
+    # Make sure all frames are the same size
+    if len(set(frame_sizes)) != 1:
+        raise ValueError("Mismatching sizes on frames")
+
+    # Write out the initial graphics descriptor (and write a dummy value), so that we can come back and fill in the
+    # correct values once we've written all the frames to the output
+    graphics_descriptor_location = fp.tell()
+    graphics_descriptor = QGFGraphicsDescriptor()
+    graphics_descriptor.frame_count = len(frame_sizes)
+    graphics_descriptor.image_size = frame_sizes[0]
+    vprint(f'{"Graphics descriptor block":26s} {fp.tell():5d}d / {fp.tell():04X}h')
+    graphics_descriptor.write(fp)
+
+    # Work out the frame offset descriptor location (and write a dummy value), so that we can come back and fill in the
+    # correct offsets once we've written all the frames to the output
+    frame_offset_location = fp.tell()
+    frame_offsets = QGFFrameOffsetDescriptorV1(graphics_descriptor.frame_count)
+    vprint(f'{"Frame offsets block":26s} {fp.tell():5d}d / {fp.tell():04X}h')
+    frame_offsets.write(fp)
+
     # Iterate over each if the input frames, writing it to the output in the process
-    _for_all_frames(_write_frame)
+    write_frame = functools.partial(_write_frame, format_=encoderinfo["qmk_format"], fp=fp, use_deltas=encoderinfo.get("use_deltas", True), use_rle=encoderinfo.get("use_rle", True), frame_offsets=frame_offsets)
+    for_all_frames(write_frame)
 
     # Go back and update the graphics descriptor now that we can determine the final file size
     graphics_descriptor.total_file_size = fp.tell()