Added reading and writing of private PNG chunks

This commit is contained in:
Andrew Murray 2020-03-11 21:40:17 +11:00
parent e2437c9b48
commit f7144c1216
2 changed files with 53 additions and 5 deletions

View File

@ -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()

View File

@ -228,15 +228,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.
@ -641,6 +646,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:
@ -657,6 +663,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)
@ -888,7 +896,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()
@ -1217,12 +1227,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
@ -1272,7 +1288,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)
@ -1290,6 +1307,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"):