Added progress callback when save_all is used

Author: radarhere

Tests/test_file_apng.py

@@ -1,3 +1,5 @@
+from io import BytesIO
+
 import pytest
 
 from PIL import Image, ImageSequence, PngImagePlugin
@@ -663,6 +665,33 @@ def test_apng_save_blend(tmp_path):
         assert im.getpixel((0, 0)) == (0, 255, 0, 255)
 
 
+def test_save_all_progress():
+    out = BytesIO()
+    progress = []
+
+    def callback(filename, frame_number, n_frames):
+        progress.append((filename, frame_number, n_frames))
+
+    Image.new("RGB", (1, 1)).save(out, "PNG", save_all=True, progress=callback)
+    assert progress == [(None, 1, 1)]
+
+    out = BytesIO()
+    progress = []
+
+    with Image.open("Tests/images/apng/single_frame.png") as im:
+        with Image.open("Tests/images/apng/delay.png") as im2:
+            im.save(out, "PNG", save_all=True, append_images=[im2], progress=callback)
+
+    assert progress == [
+        ("Tests/images/apng/single_frame.png", 1, 6),
+        ("Tests/images/apng/delay.png", 2, 6),
+        ("Tests/images/apng/delay.png", 3, 6),
+        ("Tests/images/apng/delay.png", 4, 6),
+        ("Tests/images/apng/delay.png", 5, 6),
+        ("Tests/images/apng/delay.png", 6, 6),
+    ]
+
+
 def test_seek_after_close():
     im = Image.open("Tests/images/apng/delay.png")
     im.seek(1)

Tests/test_file_gif.py

@@ -265,6 +265,33 @@ def test_roundtrip_save_all_1(tmp_path):
         assert reloaded.getpixel((0, 0)) == 255
 
 
+def test_save_all_progress():
+    out = BytesIO()
+    progress = []
+
+    def callback(filename, frame_number, n_frames):
+        progress.append((filename, frame_number, n_frames))
+
+    Image.new("RGB", (1, 1)).save(out, "GIF", save_all=True, progress=callback)
+    assert progress == [(None, 1, 1)]
+
+    out = BytesIO()
+    progress = []
+
+    with Image.open("Tests/images/hopper.gif") as im:
+        with Image.open("Tests/images/dispose_bgnd.gif") as im2:
+            im.save(out, "GIF", save_all=True, append_images=[im2], progress=callback)
+
+    assert progress == [
+        ("Tests/images/hopper.gif", 1, 6),
+        ("Tests/images/dispose_bgnd.gif", 2, 6),
+        ("Tests/images/dispose_bgnd.gif", 3, 6),
+        ("Tests/images/dispose_bgnd.gif", 4, 6),
+        ("Tests/images/dispose_bgnd.gif", 5, 6),
+        ("Tests/images/dispose_bgnd.gif", 6, 6),
+    ]
+
+
 @pytest.mark.parametrize(
     "path, mode",
     (

Tests/test_file_mpo.py

@@ -278,3 +278,28 @@ def test_save_all():
     # Test that a single frame image will not be saved as an MPO
     jpg = roundtrip(im, save_all=True)
     assert "mp" not in jpg.info
+
+
+def test_save_all_progress():
+    out = BytesIO()
+    progress = []
+
+    def callback(filename, frame_number, n_frames):
+        progress.append((filename, frame_number, n_frames))
+
+    Image.new("RGB", (1, 1)).save(out, "MPO", save_all=True, progress=callback)
+    assert progress == [(None, 1, 1)]
+
+    out = BytesIO()
+    progress = []
+
+    with Image.open("Tests/images/sugarshack.mpo") as im:
+        with Image.open("Tests/images/frozenpond.mpo") as im2:
+            im.save(out, "MPO", save_all=True, append_images=[im2], progress=callback)
+
+    assert progress == [
+        ("Tests/images/sugarshack.mpo", 1, 4),
+        ("Tests/images/sugarshack.mpo", 2, 4),
+        ("Tests/images/frozenpond.mpo", 3, 4),
+        ("Tests/images/frozenpond.mpo", 4, 4),
+    ]

Tests/test_file_pdf.py

@@ -1,8 +1,8 @@
-import io
 import os
 import os.path
 import tempfile
 import time
+from io import BytesIO
 
 import pytest
 
@@ -169,6 +169,31 @@ def im_generator(ims):
     assert os.path.getsize(outfile) > 0
 
 
+def test_save_all_progress():
+    out = BytesIO()
+    progress = []
+
+    def callback(filename, frame_number, n_frames):
+        progress.append((filename, frame_number, n_frames))
+
+    Image.new("RGB", (1, 1)).save(out, "PDF", save_all=True, progress=callback)
+    assert progress == [(None, 1, 1)]
+
+    out = BytesIO()
+    progress = []
+
+    with Image.open("Tests/images/sugarshack.mpo") as im:
+        with Image.open("Tests/images/frozenpond.mpo") as im2:
+            im.save(out, "PDF", save_all=True, append_images=[im2], progress=callback)
+
+    assert progress == [
+        ("Tests/images/sugarshack.mpo", 1, 4),
+        ("Tests/images/sugarshack.mpo", 2, 4),
+        ("Tests/images/frozenpond.mpo", 3, 4),
+        ("Tests/images/frozenpond.mpo", 4, 4),
+    ]
+
+
 def test_multiframe_normal_save(tmp_path):
     # Test saving a multiframe image without save_all
     with Image.open("Tests/images/dispose_bgnd.gif") as im:
@@ -323,12 +348,12 @@ def test_pdf_info(tmp_path):
 
 def test_pdf_append_to_bytesio():
     im = hopper("RGB")
-    f = io.BytesIO()
+    f = BytesIO()
     im.save(f, format="PDF")
     initial_size = len(f.getvalue())
     assert initial_size > 0
     im = hopper("P")
-    f = io.BytesIO(f.getvalue())
+    f = BytesIO(f.getvalue())
     im.save(f, format="PDF", append=True)
     assert len(f.getvalue()) > initial_size
 

Tests/test_file_tiff.py

@@ -658,7 +658,7 @@ def test_palette(self, mode, tmp_path):
         with Image.open(outfile) as reloaded:
             assert_image_equal(im.convert("RGB"), reloaded.convert("RGB"))
 
-    def test_tiff_save_all(self):
+    def test_save_all(self):
         mp = BytesIO()
         with Image.open("Tests/images/multipage.tiff") as im:
             im.save(mp, format="tiff", save_all=True)
@@ -688,6 +688,32 @@ def im_generator(ims):
         with Image.open(mp) as reread:
             assert reread.n_frames == 3
 
+    def test_save_all_progress(self):
+        out = BytesIO()
+        progress = []
+
+        def callback(filename, frame_number, n_frames):
+            progress.append((filename, frame_number, n_frames))
+
+        Image.new("RGB", (1, 1)).save(out, "TIFF", save_all=True, progress=callback)
+        assert progress == [(None, 1, 1)]
+
+        out = BytesIO()
+        progress = []
+
+        with Image.open("Tests/images/hopper.tif") as im:
+            with Image.open("Tests/images/multipage.tiff") as im2:
+                im.save(
+                    out, "TIFF", save_all=True, append_images=[im2], progress=callback
+                )
+
+        assert progress == [
+            ("Tests/images/hopper.tif", 1, 4),
+            ("Tests/images/multipage.tiff", 2, 4),
+            ("Tests/images/multipage.tiff", 3, 4),
+            ("Tests/images/multipage.tiff", 4, 4),
+        ]
+
     def test_saving_icc_profile(self, tmp_path):
         # Tests saving TIFF with icc_profile set.
         # At the time of writing this will only work for non-compressed tiffs

Tests/test_file_webp.py

@@ -1,7 +1,7 @@
-import io
 import re
 import sys
 import warnings
+from io import BytesIO
 
 import pytest
 
@@ -102,10 +102,10 @@ def test_write_rgb(self, tmp_path):
     def test_write_method(self, tmp_path):
         self._roundtrip(tmp_path, self.rgb_mode, 12.0, {"method": 6})
 
-        buffer_no_args = io.BytesIO()
+        buffer_no_args = BytesIO()
         hopper().save(buffer_no_args, format="WEBP")
 
-        buffer_method = io.BytesIO()
+        buffer_method = BytesIO()
         hopper().save(buffer_method, format="WEBP", method=6)
         assert buffer_no_args.getbuffer() != buffer_method.getbuffer()
 
@@ -122,6 +122,30 @@ def test_save_all(self, tmp_path):
             reloaded.seek(1)
             assert_image_similar(im2, reloaded, 1)
 
+    @skip_unless_feature("webp_anim")
+    def test_save_all_progress(self):
+        out = BytesIO()
+        progress = []
+
+        def callback(filename, frame_number, n_frames):
+            progress.append((filename, frame_number, n_frames))
+
+        Image.new("RGB", (1, 1)).save(out, "WEBP", save_all=True, progress=callback)
+        assert progress == [(None, 1, 1)]
+
+        out = BytesIO()
+        progress = []
+
+        with Image.open("Tests/images/iss634.webp") as im:
+            im2 = Image.new("RGB", im.size)
+            im.save(out, "WEBP", save_all=True, append_images=[im2], progress=callback)
+
+        expected = []
+        for i in range(42):
+            expected.append(("Tests/images/iss634.webp", i + 1, 43))
+        expected.append((None, 43, 43))
+        assert progress == expected
+
     def test_icc_profile(self, tmp_path):
         self._roundtrip(tmp_path, self.rgb_mode, 12.5, {"icc_profile": None})
         if _webp.HAVE_WEBPANIM:

src/PIL/GifImagePlugin.py

@@ -24,7 +24,6 @@
 # See the README file for information on usage and redistribution.
 #
 
-import itertools
 import math
 import os
 import subprocess
@@ -578,10 +577,17 @@ def _write_multiple_frames(im, fp, palette):
     duration = im.encoderinfo.get("duration")
     disposal = im.encoderinfo.get("disposal", im.info.get("disposal"))
 
+    progress = im.encoderinfo.get("progress")
+    imSequences = [im] + list(im.encoderinfo.get("append_images", []))
+    if progress:
+        n_frames = 0
+        for imSequence in imSequences:
+            n_frames += getattr(imSequence, "n_frames", 1)
+
     im_frames = []
     frame_count = 0
     background_im = None
-    for imSequence in itertools.chain([im], im.encoderinfo.get("append_images", [])):
+    for imSequence in imSequences:
         for im_frame in ImageSequence.Iterator(imSequence):
             # a copy is required here since seek can still mutate the image
             im_frame = _normalize_mode(im_frame.copy())
@@ -611,6 +617,10 @@ def _write_multiple_frames(im, fp, palette):
                     # This frame is identical to the previous frame
                     if encoderinfo.get("duration"):
                         previous["encoderinfo"]["duration"] += encoderinfo["duration"]
+                    if progress:
+                        progress(
+                            getattr(imSequence, "filename", None), frame_count, n_frames
+                        )
                     continue
                 if encoderinfo.get("disposal") == 2:
                     if background_im is None:
@@ -624,6 +634,8 @@ def _write_multiple_frames(im, fp, palette):
             else:
                 bbox = None
             im_frames.append({"im": im_frame, "bbox": bbox, "encoderinfo": encoderinfo})
+            if progress:
+                progress(getattr(imSequence, "filename", None), frame_count, n_frames)
 
     if len(im_frames) > 1:
         for frame_data in im_frames:

src/PIL/MpoImagePlugin.py

@@ -18,7 +18,6 @@
 # See the README file for information on usage and redistribution.
 #
 
-import itertools
 import os
 import struct
 
@@ -42,6 +41,7 @@ def _save(im, fp, filename):
 
 
 def _save_all(im, fp, filename):
+    progress = im.encoderinfo.get("progress")
     append_images = im.encoderinfo.get("append_images", [])
     if not append_images:
         try:
@@ -50,11 +50,19 @@ def _save_all(im, fp, filename):
             animated = False
         if not animated:
             _save(im, fp, filename)
+            if progress:
+                progress(getattr(im, "filename", None), 1, 1)
             return
 
     mpf_offset = 28
     offsets = []
-    for imSequence in itertools.chain([im], append_images):
+    imSequences = [im] + list(append_images)
+    if progress:
+        frame_number = 0
+        n_frames = 0
+        for imSequence in imSequences:
+            n_frames += getattr(imSequence, "n_frames", 1)
+    for imSequence in imSequences:
         for im_frame in ImageSequence.Iterator(imSequence):
             if not offsets:
                 # APP2 marker
@@ -73,6 +81,9 @@ def _save_all(im, fp, filename):
             else:
                 im_frame.save(fp, "JPEG")
                 offsets.append(fp.tell() - offsets[-1])
+            if progress:
+                frame_number += 1
+                progress(getattr(imSequence, "filename", None), frame_number, n_frames)
 
     ifd = TiffImagePlugin.ImageFileDirectory_v2()
     ifd[0xB000] = b"0100"

src/PIL/PdfImagePlugin.py

@@ -246,6 +246,7 @@ def _save(im, fp, filename, save_all=False):
     # catalog and list of pages
     existing_pdf.write_catalog()
 
+    progress = im.encoderinfo.get("progress")
     page_number = 0
     for im_sequence in ims:
         im_pages = ImageSequence.Iterator(im_sequence) if save_all else [im_sequence]
@@ -281,6 +282,10 @@ def _save(im, fp, filename, save_all=False):
             existing_pdf.write_obj(contents_refs[page_number], stream=page_contents)
 
             page_number += 1
+            if progress:
+                progress(
+                    getattr(im_sequence, "filename", None), page_number, number_of_pages
+                )
 
     #
     # trailer