Skip to content

Commit 5da88ea

Browse files
committed
Added progress callback to AVIF
1 parent 09e4df1 commit 5da88ea

7 files changed

Lines changed: 66 additions & 30 deletions

File tree

Tests/test_file_apng.py

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
from __future__ import annotations
22

3+
import os
34
from io import BytesIO
45
from pathlib import Path
56

@@ -710,11 +711,9 @@ def test_save_all_progress() -> None:
710711
out = BytesIO()
711712
progress = []
712713

713-
def callback(state):
714+
def callback(state) -> None:
714715
if state["image_filename"]:
715-
state["image_filename"] = (
716-
state["image_filename"].replace("\\", "/").split("Tests/images/")[-1]
717-
)
716+
state["image_filename"] = os.path.basename(state["image_filename"])
718717
progress.append(state)
719718

720719
Image.new("RGB", (1, 1)).save(out, "PNG", save_all=True, progress=callback)
@@ -741,7 +740,7 @@ def callback(state):
741740
expected.append(
742741
{
743742
"image_index": i,
744-
"image_filename": "apng/single_frame.png",
743+
"image_filename": "single_frame.png",
745744
"completed_frames": i + 1,
746745
"total_frames": 7,
747746
}
@@ -750,7 +749,7 @@ def callback(state):
750749
expected.append(
751750
{
752751
"image_index": 2,
753-
"image_filename": "apng/delay.png",
752+
"image_filename": "delay.png",
754753
"completed_frames": i + 3,
755754
"total_frames": 7,
756755
}

Tests/test_file_avif.py

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -217,6 +217,52 @@ def test_file_pointer_could_be_reused(self) -> None:
217217
with Image.open(blob) as im:
218218
im.load()
219219

220+
def test_save_all_progress(self) -> None:
221+
out = BytesIO()
222+
progress = []
223+
224+
def callback(state) -> None:
225+
if state["image_filename"]:
226+
state["image_filename"] = os.path.basename(state["image_filename"])
227+
progress.append(state)
228+
229+
Image.new("RGB", (1, 1)).save(out, "AVIF", save_all=True, progress=callback)
230+
assert progress == [
231+
{
232+
"image_index": 0,
233+
"image_filename": None,
234+
"completed_frames": 1,
235+
"total_frames": 1,
236+
}
237+
]
238+
239+
out = BytesIO()
240+
progress = []
241+
242+
with Image.open("Tests/images/avif/star.avifs") as im:
243+
im2 = Image.new(im.mode, im.size)
244+
im.save(out, "AVIF", save_all=True, append_images=[im2], progress=callback)
245+
246+
expected = []
247+
for i in range(5):
248+
expected.append(
249+
{
250+
"image_index": 0,
251+
"image_filename": "star.avifs",
252+
"completed_frames": i + 1,
253+
"total_frames": 6,
254+
}
255+
)
256+
expected.append(
257+
{
258+
"image_index": 1,
259+
"image_filename": None,
260+
"completed_frames": 6,
261+
"total_frames": 6,
262+
}
263+
)
264+
assert progress == expected
265+
220266
def test_background_from_gif(self, tmp_path: Path) -> None:
221267
with Image.open("Tests/images/chi.gif") as im:
222268
original_value = im.convert("RGB").getpixel((1, 1))

Tests/test_file_gif.py

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
from __future__ import annotations
22

3+
import os
34
import warnings
45
from collections.abc import Generator
56
from io import BytesIO
@@ -296,15 +297,13 @@ def test_roundtrip_save_all_1(tmp_path: Path) -> None:
296297
assert reloaded.getpixel((0, 0)) == 255
297298

298299

299-
def test_save_all_progress():
300+
def test_save_all_progress() -> None:
300301
out = BytesIO()
301302
progress = []
302303

303-
def callback(state):
304+
def callback(state) -> None:
304305
if state["image_filename"]:
305-
state["image_filename"] = (
306-
state["image_filename"].replace("\\", "/").split("Tests/images/")[-1]
307-
)
306+
state["image_filename"] = os.path.basename(state["image_filename"])
308307
progress.append(state)
309308

310309
Image.new("RGB", (1, 1)).save(out, "GIF", save_all=True, progress=callback)
@@ -324,7 +323,7 @@ def callback(state):
324323
im = Image.new("RGB", im2.size)
325324
im.save(out, "GIF", save_all=True, append_images=[im2], progress=callback)
326325

327-
expected = [
326+
expected: list[dict[str, int | str | None]] = [
328327
{
329328
"image_index": 0,
330329
"image_filename": None,

Tests/test_file_pdf.py

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -178,11 +178,9 @@ def test_save_all_progress() -> None:
178178
out = BytesIO()
179179
progress = []
180180

181-
def callback(state):
181+
def callback(state) -> None:
182182
if state["image_filename"]:
183-
state["image_filename"] = (
184-
state["image_filename"].replace("\\", "/").split("Tests/images/")[-1]
185-
)
183+
state["image_filename"] = os.path.basename(state["image_filename"])
186184
progress.append(state)
187185

188186
Image.new("RGB", (1, 1)).save(out, "PDF", save_all=True, progress=callback)

Tests/test_file_tiff.py

Lines changed: 2 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -797,13 +797,9 @@ def test_save_all_progress(self) -> None:
797797
out = BytesIO()
798798
progress = []
799799

800-
def callback(state):
800+
def callback(state) -> None:
801801
if state["image_filename"]:
802-
state["image_filename"] = (
803-
state["image_filename"]
804-
.replace("\\", "/")
805-
.split("Tests/images/")[-1]
806-
)
802+
state["image_filename"] = os.path.basename(state["image_filename"])
807803
progress.append(state)
808804

809805
Image.new("RGB", (1, 1)).save(out, "TIFF", save_all=True, progress=callback)

Tests/test_file_webp.py

Lines changed: 3 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
from __future__ import annotations
22

3+
import os
34
import re
45
import sys
56
import warnings
@@ -129,18 +130,13 @@ def test_unsupported_image_mode(self) -> None:
129130
with pytest.raises(ValueError):
130131
_webp.WebPEncode(im.getim(), False, 0, 0, "", 4, 0, b"", "")
131132

132-
@skip_unless_feature("webp_anim")
133133
def test_save_all_progress(self) -> None:
134134
out = BytesIO()
135135
progress = []
136136

137-
def callback(state):
137+
def callback(state) -> None:
138138
if state["image_filename"]:
139-
state["image_filename"] = (
140-
state["image_filename"]
141-
.replace("\\", "/")
142-
.split("Tests/images/")[-1]
143-
)
139+
state["image_filename"] = os.path.basename(state["image_filename"])
144140
progress.append(state)
145141

146142
Image.new("RGB", (1, 1)).save(out, "WEBP", save_all=True, progress=callback)

src/PIL/AvifImagePlugin.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -236,8 +236,9 @@ def _save(
236236
frame_duration = 0
237237
cur_idx = im.tell()
238238
is_single_frame = total == 1
239+
progress = info.get("progress")
239240
try:
240-
for ims in [im] + append_images:
241+
for i, ims in enumerate([im] + append_images):
241242
# Get number of frames in this image
242243
nfr = getattr(ims, "n_frames", 1)
243244

@@ -268,6 +269,7 @@ def _save(
268269

269270
# Update frame index
270271
frame_idx += 1
272+
im._save_all_progress(progress, ims, i, frame_idx, total)
271273

272274
if not save_all:
273275
break

0 commit comments

Comments
 (0)