Merge pull request #4292 from radarhere/private_png_chunks

Added reading and writing of private PNG chunks

Tests/test_file_png.py

@@ -564,6 +564,28 @@ class TestFilePng:
         chunks = PngImagePlugin.getchunks(im)
         assert len(chunks) == 3
 
+    def test_read_private_chunks(self):
+        im = Image.open("Tests/images/exif.png")
+        assert im.private_chunks == [(b"orNT", b"\x01")]
+
+    def test_roundtrip_private_chunk(self):
+        # Check private chunk roundtripping
+
+        with Image.open(TEST_PNG_FILE) as im:
+            info = PngImagePlugin.PngInfo()
+            info.add(b"prIV", b"VALUE")
+            info.add(b"atEC", b"VALUE2")
+            info.add(b"prIV", b"VALUE3", True)
+
+            im = roundtrip(im, pnginfo=info)
+        assert im.private_chunks == [(b"prIV", b"VALUE"), (b"atEC", b"VALUE2")]
+        im.load()
+        assert im.private_chunks == [
+            (b"prIV", b"VALUE"),
+            (b"atEC", b"VALUE2"),
+            (b"prIV", b"VALUE3", True),
+        ]
+
     def test_textual_chunks_after_idat(self):
         with Image.open("Tests/images/hopper.png") as im:
             assert "comment" in im.text.keys()

docs/handbook/image-file-formats.rst

@@ -535,7 +535,7 @@ The :py:meth:`~PIL.Image.Image.save` method supports the following options:
     A tuple of two numbers corresponding to the desired dpi in each direction.
 
 **pnginfo**
-    A :py:class:`PIL.PngImagePlugin.PngInfo` instance containing text tags.
+    A :py:class:`PIL.PngImagePlugin.PngInfo` instance containing chunks.
 
 **compress_level**
     ZLIB compression level, a number between 0 and 9: 1 gives best speed,

src/PIL/PngImagePlugin.py

@@ -266,15 +266,20 @@ class PngInfo:
     def __init__(self):
         self.chunks = []
 
-    def add(self, cid, data):
+    def add(self, cid, data, after_idat=False):
         """Appends an arbitrary chunk. Use with caution.
 
         :param cid: a byte string, 4 bytes long.
         :param data: a byte string of the encoded data
+        :param after_idat: for use with private chunks. Whether the chunk
+                           should be written after IDAT
 
         """
 
-        self.chunks.append((cid, data))
+        chunk = [cid, data]
+        if after_idat:
+            chunk.append(True)
+        self.chunks.append(tuple(chunk))
 
     def add_itxt(self, key, value, lang="", tkey="", zip=False):
         """Appends an iTXt chunk.
@@ -676,6 +681,7 @@ class PngImageFile(ImageFile.ImageFile):
         #
         # Parse headers up to the first IDAT or fDAT chunk
 
+        self.private_chunks = []
         self.png = PngStream(self.fp)
 
         while True:
@@ -692,6 +698,8 @@ class PngImageFile(ImageFile.ImageFile):
             except AttributeError:
                 logger.debug("%r %s %s (unknown)", cid, pos, length)
                 s = ImageFile._safe_read(self.fp, length)
+                if cid[1:2].islower():
+                    self.private_chunks.append((cid, s))
 
             self.png.crc(cid, s)
 
@@ -923,7 +931,9 @@ class PngImageFile(ImageFile.ImageFile):
                 ImageFile._safe_read(self.fp, length)
             except AttributeError:
                 logger.debug("%r %s %s (unknown)", cid, pos, length)
-                ImageFile._safe_read(self.fp, length)
+                s = ImageFile._safe_read(self.fp, length)
+                if cid[1:2].islower():
+                    self.private_chunks.append((cid, s, True))
         self._text = self.png.im_text
         if not self.is_animated:
             self.png.close()
@@ -1248,12 +1258,18 @@ def _save(im, fp, filename, chunk=putchunk, save_all=False):
     info = im.encoderinfo.get("pnginfo")
     if info:
         chunks_multiple_allowed = [b"sPLT", b"iTXt", b"tEXt", b"zTXt"]
-        for cid, data in info.chunks:
+        for info_chunk in info.chunks:
+            cid, data = info_chunk[:2]
             if cid in chunks:
                 chunks.remove(cid)
                 chunk(fp, cid, data)
             elif cid in chunks_multiple_allowed:
                 chunk(fp, cid, data)
+            elif cid[1:2].islower():
+                # Private chunk
+                after_idat = info_chunk[2:3]
+                if not after_idat:
+                    chunk(fp, cid, data)
 
     if im.mode == "P":
         palette_byte_number = (2 ** bits) * 3
@@ -1303,7 +1319,8 @@ def _save(im, fp, filename, chunk=putchunk, save_all=False):
 
     if info:
         chunks = [b"bKGD", b"hIST"]
-        for cid, data in info.chunks:
+        for info_chunk in info.chunks:
+            cid, data = info_chunk[:2]
             if cid in chunks:
                 chunks.remove(cid)
                 chunk(fp, cid, data)
@@ -1321,6 +1338,15 @@ def _save(im, fp, filename, chunk=putchunk, save_all=False):
     else:
         ImageFile._save(im, _idat(fp, chunk), [("zip", (0, 0) + im.size, 0, rawmode)])
 
+    if info:
+        for info_chunk in info.chunks:
+            cid, data = info_chunk[:2]
+            if cid[1:2].islower():
+                # Private chunk
+                after_idat = info_chunk[2:3]
+                if after_idat:
+                    chunk(fp, cid, data)
+
     chunk(fp, b"IEND", b"")
 
     if hasattr(fp, "flush"):