diff --git a/Tests/test_file_png.py b/Tests/test_file_png.py index 3b4ffff3ef0..38d7033110c 100644 --- a/Tests/test_file_png.py +++ b/Tests/test_file_png.py @@ -777,6 +777,16 @@ class MyStdOut: mystdout = mystdout.buffer with Image.open(mystdout) as reloaded: assert_image_equal_tofile(reloaded, TEST_PNG_FILE) + + def test_end_truncated_filed(self): + ImageFile.LOAD_TRUNCATED_IMAGES = True + try: + img = Image.open("Tests/images/end_trunc_file.png") + img.load() + assert img is not None + finally: + ImageFile.LOAD_TRUNCATED_IMAGES = False + @pytest.mark.skipif(is_win32(), reason="Requires Unix or macOS") @@ -786,20 +796,18 @@ class TestTruncatedPngPLeaks(PillowLeakTestCase): iterations = 100 # Leak is 56k/iteration, this will leak 5.6megs def test_leak_load(self): - files = ['Tests/images/hopper.png', 'Tests/images/end_trunc_file.png'] - for file in files: - with open(file, "rb") as f: - DATA = BytesIO(f.read(16 * 1024)) + with open("Tests/images/hopper.png", "rb") as f: + DATA = BytesIO(f.read(16 * 1024)) - ImageFile.LOAD_TRUNCATED_IMAGES = True + ImageFile.LOAD_TRUNCATED_IMAGES = True + with Image.open(DATA) as im: + im.load() + + def core(): with Image.open(DATA) as im: im.load() - def core(): - with Image.open(DATA) as im: - im.load() - - try: - self._test_leak(core) - finally: - ImageFile.LOAD_TRUNCATED_IMAGES = False + try: + self._test_leak(core) + finally: + ImageFile.LOAD_TRUNCATED_IMAGES = False diff --git a/src/PIL/PngImagePlugin.py b/src/PIL/PngImagePlugin.py index 7797b8e6174..f51773a348e 100644 --- a/src/PIL/PngImagePlugin.py +++ b/src/PIL/PngImagePlugin.py @@ -984,9 +984,7 @@ def load_end(self): try: ImageFile._safe_read(self.fp, length) except OSError as e: - if ( - str(e) == "Truncated File Read" - ) and ImageFile.LOAD_TRUNCATED_IMAGES: + if ImageFile.LOAD_TRUNCATED_IMAGES: logger.debug("%r %s %s (Truncated File Read)", cid, pos, length) break else: