Run encoder cleanup method after errors as well

This commit is contained in:
Andrew Murray 2022-02-28 14:12:56 +11:00
parent 2f2b48dc2c
commit a606fd85a3
2 changed files with 35 additions and 26 deletions

View File

@ -200,6 +200,9 @@ class MockPyEncoder(ImageFile.PyEncoder):
def encode(self, buffer): def encode(self, buffer):
return 1, 1, b"" return 1, 1, b""
def cleanup(self):
self.cleanup_called = True
xoff, yoff, xsize, ysize = 10, 20, 100, 100 xoff, yoff, xsize, ysize = 10, 20, 100, 100
@ -327,10 +330,12 @@ class TestPyEncoder(CodecsTest):
im = MockImageFile(buf) im = MockImageFile(buf)
fp = BytesIO() fp = BytesIO()
self.encoder.cleanup_called = False
with pytest.raises(ValueError): with pytest.raises(ValueError):
ImageFile._save( ImageFile._save(
im, fp, [("MOCK", (xoff, yoff, -10, yoff + ysize), 0, "RGB")] im, fp, [("MOCK", (xoff, yoff, -10, yoff + ysize), 0, "RGB")]
) )
assert self.encoder.cleanup_called
with pytest.raises(ValueError): with pytest.raises(ValueError):
ImageFile._save( ImageFile._save(

View File

@ -503,36 +503,40 @@ def _save(im, fp, tile, bufsize=0):
# compress to Python file-compatible object # compress to Python file-compatible object
for e, b, o, a in tile: for e, b, o, a in tile:
e = Image._getencoder(im.mode, e, a, im.encoderconfig) e = Image._getencoder(im.mode, e, a, im.encoderconfig)
if o > 0: try:
fp.seek(o) if o > 0:
e.setimage(im.im, b) fp.seek(o)
if e.pushes_fd: e.setimage(im.im, b)
e.setfd(fp) if e.pushes_fd:
l, s = e.encode_to_pyfd() e.setfd(fp)
else: l, s = e.encode_to_pyfd()
while True: else:
l, s, d = e.encode(bufsize) while True:
fp.write(d) l, s, d = e.encode(bufsize)
if s: fp.write(d)
break if s:
if s < 0: break
raise OSError(f"encoder error {s} when writing image file") from exc if s < 0:
e.cleanup() raise OSError(f"encoder error {s} when writing image file") from exc
finally:
e.cleanup()
else: else:
# slight speedup: compress to real file object # slight speedup: compress to real file object
for e, b, o, a in tile: for e, b, o, a in tile:
e = Image._getencoder(im.mode, e, a, im.encoderconfig) e = Image._getencoder(im.mode, e, a, im.encoderconfig)
if o > 0: try:
fp.seek(o) if o > 0:
e.setimage(im.im, b) fp.seek(o)
if e.pushes_fd: e.setimage(im.im, b)
e.setfd(fp) if e.pushes_fd:
l, s = e.encode_to_pyfd() e.setfd(fp)
else: l, s = e.encode_to_pyfd()
s = e.encode_to_file(fh, bufsize) else:
if s < 0: s = e.encode_to_file(fh, bufsize)
raise OSError(f"encoder error {s} when writing image file") if s < 0:
e.cleanup() raise OSError(f"encoder error {s} when writing image file")
finally:
e.cleanup()
if hasattr(fp, "flush"): if hasattr(fp, "flush"):
fp.flush() fp.flush()