Pillow/Tests/check_j2k_leaks.py

47 lines
1.4 KiB
Python
Raw Normal View History

import unittest
from io import BytesIO
from PIL import Image
from .helper import PillowTestCase, is_win32
# Limits for testing the leak
2019-06-13 18:53:42 +03:00
mem_limit = 1024 * 1048576
stack_size = 8 * 1048576
iterations = int((mem_limit / stack_size) * 2)
codecs = dir(Image.core)
test_file = "Tests/images/rgb_trns_ycbc.jp2"
2019-09-25 12:46:54 +03:00
@unittest.skipIf(is_win32(), "requires Unix or macOS")
class TestJpegLeaks(PillowTestCase):
def setUp(self):
if "jpeg2k_encoder" not in codecs or "jpeg2k_decoder" not in codecs:
2019-06-13 18:53:42 +03:00
self.skipTest("JPEG 2000 support not available")
def test_leak_load(self):
from resource import setrlimit, RLIMIT_AS, RLIMIT_STACK
2019-06-13 18:53:42 +03:00
setrlimit(RLIMIT_STACK, (stack_size, stack_size))
setrlimit(RLIMIT_AS, (mem_limit, mem_limit))
2015-04-08 14:12:37 +03:00
for _ in range(iterations):
with Image.open(test_file) as im:
im.load()
def test_leak_save(self):
from resource import setrlimit, RLIMIT_AS, RLIMIT_STACK
2019-06-13 18:53:42 +03:00
setrlimit(RLIMIT_STACK, (stack_size, stack_size))
setrlimit(RLIMIT_AS, (mem_limit, mem_limit))
2015-04-08 14:12:37 +03:00
for _ in range(iterations):
with Image.open(test_file) as im:
im.load()
test_output = BytesIO()
im.save(test_output, "JPEG2000")
test_output.seek(0)
2015-04-08 14:12:37 +03:00
test_output.read()
2019-06-13 18:53:42 +03:00
if __name__ == "__main__":
unittest.main()