Support saving multiple MPO frames

Tests/test_file_mpo.py

@@ -5,15 +5,19 @@
 
 from PIL import Image
 
-from .helper import assert_image_similar, is_pypy, skip_unless_feature
+from .helper import (
+    assert_image_equal,
+    assert_image_similar,
+    is_pypy,
+    skip_unless_feature,
+)
 
 test_files = ["Tests/images/sugarshack.mpo", "Tests/images/frozenpond.mpo"]
 
 pytestmark = skip_unless_feature("jpg")
 
 
-def frame_roundtrip(im, **options):
-    # Note that for now, there is no MPO saving functionality
+def roundtrip(im, **options):
     out = BytesIO()
     im.save(out, "MPO", **options)
     test_bytes = out.tell()
@@ -237,13 +241,38 @@ def test_image_grab():
 
 
 def test_save():
-    # Note that only individual frames can be saved at present
     for test_file in test_files:
         with Image.open(test_file) as im:
             assert im.tell() == 0
-            jpg0 = frame_roundtrip(im)
+            jpg0 = roundtrip(im)
             assert_image_similar(im, jpg0, 30)
             im.seek(1)
             assert im.tell() == 1
-            jpg1 = frame_roundtrip(im)
+            jpg1 = roundtrip(im)
             assert_image_similar(im, jpg1, 30)
+
+
+def test_save_all():
+    for test_file in test_files:
+        with Image.open(test_file) as im:
+            im_reloaded = roundtrip(im, save_all=True)
+
+            im.seek(0)
+            assert_image_similar(im, im_reloaded, 30)
+
+            im.seek(1)
+            im_reloaded.seek(1)
+            assert_image_similar(im, im_reloaded, 30)
+
+    im = Image.new("RGB", (1, 1))
+    im2 = Image.new("RGB", (1, 1), "#f00")
+    im_reloaded = roundtrip(im, save_all=True, append_images=[im2])
+
+    assert_image_equal(im, im_reloaded)
+
+    im_reloaded.seek(1)
+    assert_image_similar(im2, im_reloaded, 1)
+
+    # 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

docs/handbook/image-file-formats.rst

@@ -1209,6 +1209,17 @@ image when first opened. The :py:meth:`~PIL.Image.Image.seek` and :py:meth:`~PIL
 methods may be used to read other pictures from the file. The pictures are
 zero-indexed and random access is supported.
 
+When calling :py:meth:`~PIL.Image.Image.save` to write an MPO file, by default
+only the first frame of a multiframe image will be saved. If the ``save_all``
+argument is present and true, then all frames will be saved, and the following
+option will also be available.
+
+**append_images**
+    A list of images to append as additional pictures. Each of the
+    images in the list can be single or multiframe images.
+
+    .. versionadded:: 9.3.0
+
 PCD
 ^^^
 

docs/releasenotes/9.3.0.rst

@@ -0,0 +1,59 @@
+9.3.0
+-----
+
+Backwards Incompatible Changes
+==============================
+
+TODO
+^^^^
+
+Deprecations
+============
+
+TODO
+^^^^
+
+TODO
+
+API Changes
+===========
+
+TODO
+^^^^
+
+TODO
+
+API Additions
+=============
+
+Saving multiple MPO frames
+^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Multiple MPO frames can now be saved. Using the ``save_all`` argument, all of
+an image's frames will be saved to file::
+
+    from PIL import Image
+    im = Image.open("frozenpond.mpo")
+    im.save(out, save_all=True)
+
+Additional images can also be appended when saving, by combining the
+``save_all`` argument with the ``append_images`` argument::
+
+    im.save(out, save_all=True, append_images=[im1, im2, ...])
+
+
+Security
+========
+
+TODO
+^^^^
+
+TODO
+
+Other Changes
+=============
+
+TODO
+^^^^
+
+TODO

docs/releasenotes/index.rst

@@ -14,6 +14,7 @@ expected to be backported to earlier versions.
 .. toctree::
   :maxdepth: 2
 
+  9.3.0
   9.2.0
   9.1.1
   9.1.0

src/PIL/JpegImagePlugin.py

@@ -711,7 +711,7 @@ def validate_qtables(qtables):
         qtables = getattr(im, "quantization", None)
     qtables = validate_qtables(qtables)
 
-    extra = b""
+    extra = info.get("extra", b"")
 
     icc_profile = info.get("icc_profile")
     if icc_profile:

src/PIL/MpoImagePlugin.py

@@ -18,16 +18,66 @@
 # See the README file for information on usage and redistribution.
 #
 
-from . import Image, ImageFile, JpegImagePlugin
+import itertools
+import os
+import struct
+
+from . import Image, ImageFile, ImageSequence, JpegImagePlugin, TiffImagePlugin
 from ._binary import i16be as i16
+from ._binary import o32le
 
 # def _accept(prefix):
 #     return JpegImagePlugin._accept(prefix)
 
 
 def _save(im, fp, filename):
-    # Note that we can only save the current frame at present
-    return JpegImagePlugin._save(im, fp, filename)
+    JpegImagePlugin._save(im, fp, filename)
+
+
+def _save_all(im, fp, filename):
+    append_images = im.encoderinfo.get("append_images", [])
+    if not append_images:
+        try:
+            animated = im.is_animated
+        except AttributeError:
+            animated = False
+        if not animated:
+            _save(im, fp, filename)
+            return
+
+    offsets = []
+    for imSequence in itertools.chain([im], append_images):
+        for im_frame in ImageSequence.Iterator(imSequence):
+            if not offsets:
+                # APP2 marker
+                im.encoderinfo["extra"] = (
+                    b"\xFF\xE2" + struct.pack(">H", 6 + 70) + b"MPF\0" + b" " * 70
+                )
+                JpegImagePlugin._save(im_frame, fp, filename)
+                offsets.append(fp.tell())
+            else:
+                im_frame.save(fp, "JPEG")
+                offsets.append(fp.tell() - offsets[-1])
+
+    ifd = TiffImagePlugin.ImageFileDirectory_v2()
+    ifd[0xB001] = len(offsets)
+
+    mpentries = b""
+    data_offset = 0
+    for i, size in enumerate(offsets):
+        if i == 0:
+            mptype = 0x030000  # Baseline MP Primary Image
+        else:
+            mptype = 0x000000  # Undefined
+        mpentries += struct.pack("<LLLHH", mptype, size, data_offset, 0, 0)
+        if i == 0:
+            data_offset -= 28
+        data_offset += size
+    ifd[0xB002] = mpentries
+
+    fp.seek(28)
+    fp.write(b"II\x2A\x00" + o32le(8) + ifd.tobytes(8))
+    fp.seek(0, os.SEEK_END)
 
 
 ##
@@ -124,6 +174,7 @@ def adopt(jpeg_instance, mpheader=None):
 # Image.register_open(MpoImageFile.format,
 #                     JpegImagePlugin.jpeg_factory, _accept)
 Image.register_save(MpoImageFile.format, _save)
+Image.register_save_all(MpoImageFile.format, _save_all)
 
 Image.register_extension(MpoImageFile.format, ".mpo")