Skip to content

Commit

Permalink
Fix test and code.
Browse files Browse the repository at this point in the history
  • Loading branch information
lajiyuan committed Jan 12, 2024
1 parent 23fd509 commit 388a3f6
Show file tree
Hide file tree
Showing 2 changed files with 22 additions and 16 deletions.
34 changes: 21 additions & 13 deletions Tests/test_file_png.py
Expand Up @@ -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")
Expand All @@ -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
4 changes: 1 addition & 3 deletions src/PIL/PngImagePlugin.py
Expand Up @@ -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:
Expand Down

0 comments on commit 388a3f6

Please sign in to comment.