Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Ensure cleanup() is called for PyEncoders #6096

Merged
merged 3 commits into from Mar 6, 2022
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
5 changes: 5 additions & 0 deletions Tests/test_imagefile.py
Expand Up @@ -200,6 +200,9 @@ class MockPyEncoder(ImageFile.PyEncoder):
def encode(self, buffer):
return 1, 1, b""

def cleanup(self):
self.cleanup_called = True


xoff, yoff, xsize, ysize = 10, 20, 100, 100

Expand Down Expand Up @@ -327,10 +330,12 @@ def test_negsize(self):
im = MockImageFile(buf)

fp = BytesIO()
self.encoder.cleanup_called = False
with pytest.raises(ValueError):
ImageFile._save(
im, fp, [("MOCK", (xoff, yoff, -10, yoff + ysize), 0, "RGB")]
)
assert self.encoder.cleanup_called

with pytest.raises(ValueError):
ImageFile._save(
Expand Down
57 changes: 25 additions & 32 deletions src/PIL/ImageFile.py
Expand Up @@ -223,11 +223,11 @@ def load(self):
)
]
for decoder_name, extents, offset, args in self.tile:
seek(offset)
decoder = Image._getdecoder(
self.mode, decoder_name, args, self.decoderconfig
)
try:
seek(offset)
decoder.setimage(self.im, extents)
if decoder.pulls_fd:
decoder.setfd(self.fp)
Expand Down Expand Up @@ -499,40 +499,33 @@ def _save(im, fp, tile, bufsize=0):
try:
fh = fp.fileno()
fp.flush()
except (AttributeError, io.UnsupportedOperation) as exc:
# compress to Python file-compatible object
for e, b, o, a in tile:
e = Image._getencoder(im.mode, e, a, im.encoderconfig)
if o > 0:
fp.seek(o)
e.setimage(im.im, b)
if e.pushes_fd:
e.setfd(fp)
l, s = e.encode_to_pyfd()
exc = None
except (AttributeError, io.UnsupportedOperation) as e:
exc = e
for e, b, o, a in tile:
if o > 0:
fp.seek(o)
encoder = Image._getencoder(im.mode, e, a, im.encoderconfig)
try:
encoder.setimage(im.im, b)
if encoder.pushes_fd:
encoder.setfd(fp)
l, s = encoder.encode_to_pyfd()
else:
while True:
l, s, d = e.encode(bufsize)
fp.write(d)
if s:
break
if exc:
# compress to Python file-compatible object
while True:
l, s, d = encoder.encode(bufsize)
fp.write(d)
if s:
break
else:
# slight speedup: compress to real file object
s = encoder.encode_to_file(fh, bufsize)
if s < 0:
raise OSError(f"encoder error {s} when writing image file") from exc
e.cleanup()
else:
# slight speedup: compress to real file object
for e, b, o, a in tile:
e = Image._getencoder(im.mode, e, a, im.encoderconfig)
if o > 0:
fp.seek(o)
e.setimage(im.im, b)
if e.pushes_fd:
e.setfd(fp)
l, s = e.encode_to_pyfd()
else:
s = e.encode_to_file(fh, bufsize)
if s < 0:
raise OSError(f"encoder error {s} when writing image file")
e.cleanup()
finally:
encoder.cleanup()
if hasattr(fp, "flush"):
fp.flush()

Expand Down