Skip to content

Commit

Permalink
Added progress callback when save_all is used
Browse files Browse the repository at this point in the history
  • Loading branch information
radarhere committed Oct 2, 2023
1 parent aaa7587 commit 6850465
Show file tree
Hide file tree
Showing 12 changed files with 232 additions and 19 deletions.
32 changes: 32 additions & 0 deletions Tests/test_file_apng.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
from io import BytesIO

import pytest

from PIL import Image, ImageSequence, PngImagePlugin
Expand Down Expand Up @@ -663,6 +665,36 @@ 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=[im, im2], progress=callback
)

assert progress == [
("Tests/images/apng/single_frame.png", 1, 7),
("Tests/images/apng/single_frame.png", 2, 7),
("Tests/images/apng/delay.png", 3, 7),
("Tests/images/apng/delay.png", 4, 7),
("Tests/images/apng/delay.png", 5, 7),
("Tests/images/apng/delay.png", 6, 7),
("Tests/images/apng/delay.png", 7, 7),
]


def test_seek_after_close():
im = Image.open("Tests/images/apng/delay.png")
im.seek(1)
Expand Down
23 changes: 23 additions & 0 deletions Tests/test_file_gif.py
Original file line number Diff line number Diff line change
Expand Up @@ -265,6 +265,29 @@ 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/chi.gif") as im2:
im.save(out, "GIF", save_all=True, append_images=[im2], progress=callback)

expected = [("Tests/images/hopper.gif", 1, 32)]
for i in range(31):
expected.append(("Tests/images/chi.gif", i + 2, 32))
assert progress == expected


@pytest.mark.parametrize(
"path, mode",
(
Expand Down
25 changes: 25 additions & 0 deletions Tests/test_file_mpo.py
Original file line number Diff line number Diff line change
Expand Up @@ -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),
]
31 changes: 28 additions & 3 deletions Tests/test_file_pdf.py
Original file line number Diff line number Diff line change
@@ -1,8 +1,8 @@
import io
import os
import os.path
import tempfile
import time
from io import BytesIO

import pytest

Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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

Expand Down
28 changes: 27 additions & 1 deletion Tests/test_file_tiff.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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
Expand Down
30 changes: 27 additions & 3 deletions Tests/test_file_webp.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
import io
import re
import sys
import warnings
from io import BytesIO

import pytest

Expand Down Expand Up @@ -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()

Expand All @@ -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:
Expand Down
16 changes: 14 additions & 2 deletions src/PIL/GifImagePlugin.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,6 @@
# See the README file for information on usage and redistribution.
#

import itertools
import math
import os
import subprocess
Expand Down Expand Up @@ -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())
Expand Down Expand Up @@ -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:
Expand All @@ -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:
Expand Down
15 changes: 13 additions & 2 deletions src/PIL/MpoImagePlugin.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,6 @@
# See the README file for information on usage and redistribution.
#

import itertools
import os
import struct

Expand All @@ -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:
Expand All @@ -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
Expand All @@ -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"
Expand Down
5 changes: 5 additions & 0 deletions src/PIL/PdfImagePlugin.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down Expand Up @@ -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
Expand Down
Loading

0 comments on commit 6850465

Please sign in to comment.