From 4ac4cf348bd80ec49e555138ec2350526182d0e0 Mon Sep 17 00:00:00 2001
From: Alexander Piskun <13381981+bigcat88@users.noreply.github.com>
Date: Fri, 29 Nov 2024 09:37:31 +0300
Subject: [PATCH] using `encoderinfo` for saving also for the non-primary
 images (#313)

Signed-off-by: Alexander Piskun <bigcat88@icloud.com>
---
 pillow_heif/as_plugin.py |  3 +--
 tests/write_test.py      | 39 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 40 insertions(+), 2 deletions(-)

diff --git a/pillow_heif/as_plugin.py b/pillow_heif/as_plugin.py
index a1c85409..324d6430 100644
--- a/pillow_heif/as_plugin.py
+++ b/pillow_heif/as_plugin.py
@@ -272,8 +272,7 @@ def _pil_encode_image(ctx: CtxEncode, img: Image.Image, primary: bool, **kwargs)
     _info = img.info.copy()
     _info["exif"] = _exif_from_pillow(img)
     _info["xmp"] = _xmp_from_pillow(img)
-    if primary:
-        _info.update(**kwargs)
+    _info.update(**kwargs)
     _info["primary"] = primary
     if img.mode == "YCbCr":
         ctx.add_image_ycbcr(img, image_orientation=_get_orientation_for_encoder(_info), **_info)
diff --git a/tests/write_test.py b/tests/write_test.py
index 407ba172..043a518a 100644
--- a/tests/write_test.py
+++ b/tests/write_test.py
@@ -1,5 +1,7 @@
 import builtins
+import math
 import os
+import zlib
 from gc import collect
 from io import SEEK_END, BytesIO
 from pathlib import Path
@@ -583,6 +585,43 @@ def test_lossless_encoding_rgba(save_format):
     helpers.assert_image_equal(im_rgb, Image.open(buf))
 
 
+@pytest.mark.parametrize("save_format", ("HEIF", "AVIF"))
+def test_lossless_encoding_non_primary_image(save_format):
+    def encode_to_image(data):
+        binary_data = zlib.compress(data)
+        side_length = math.ceil(math.sqrt((len(binary_data) + 3) // 4))
+        total_pixels = side_length * side_length
+        padded_binary_data = binary_data + b"\x00" * (total_pixels * 4 - len(binary_data))
+        binary_array = [tuple(padded_binary_data[i : i + 4]) for i in range(0, len(padded_binary_data), 4)]
+        image_data = [binary_array[i : i + side_length] for i in range(0, len(binary_array), side_length)]
+        im = Image.new("RGBA", (side_length, side_length))
+        im.putdata([pixel for row in image_data for pixel in row])
+        return im
+
+    im_rgb = helpers.gradient_rgb()
+    im_rgb.encoderinfo = {"matrix_coefficients": 0}
+    binary_data_to_encode = b"Some binary data to encode"
+    second_image = encode_to_image(binary_data_to_encode)
+    second_image.encoderinfo = {"matrix_coefficients": 0}
+    buf = BytesIO()
+    im_rgb.save(
+        buf,
+        format=save_format,
+        quality=-1,
+        chroma=444,
+        matrix_coefficients=0,
+        save_all=True,
+        append_images=[second_image],
+    )
+    im_out = Image.open(buf)
+    helpers.assert_image_equal(im_rgb, im_out)
+    bin_data = ImageSequence.Iterator(im_out)[1].copy()
+    flattened_binary_array = list(bin_data.getdata())
+    binary_data_with_padding = b"".join(bytes(pixel) for pixel in flattened_binary_array)
+    z = binary_data_with_padding.rstrip(b"\x00")
+    assert zlib.decompress(z) == binary_data_to_encode
+
+
 def test_input_chroma_value():
     im = Image.open(Path("images/heif_other/RGB_8_chroma444.heif"))
     assert im.info["chroma"] == 444