diff --git a/tests/__init__.py b/tests/__init__.py index d81f2c94c..0da6f6548 100644 --- a/tests/__init__.py +++ b/tests/__init__.py @@ -143,3 +143,25 @@ def test_csv_consistency(): # Ensure the urls are unique assert len(pdfs) == len({pdf["url"] for pdf in pdfs}) + + +class PILContext: + """Allow changing the PIL/Pillow configuration for some limited scope.""" + + def __init__(self): + self._saved_load_truncated_images = False + + def __enter__(self): + # Allow loading incomplete images. + from PIL import ImageFile + self._saved_load_truncated_images = ImageFile.LOAD_TRUNCATED_IMAGES + ImageFile.LOAD_TRUNCATED_IMAGES = True + return self + + def __exit__(self, type_, value, traceback): + from PIL import ImageFile + ImageFile.LOAD_TRUNCATED_IMAGES = self._saved_load_truncated_images + if type_: + # Error. + return + return True diff --git a/tests/test_filters.py b/tests/test_filters.py index a2292d967..d3980be0b 100644 --- a/tests/test_filters.py +++ b/tests/test_filters.py @@ -21,7 +21,7 @@ ) from pypdf.generic import ArrayObject, DictionaryObject, NameObject, NumberObject -from . import get_data_from_url +from . import PILContext, get_data_from_url from .test_encryption import HAS_AES from .test_images import image_similarity @@ -371,13 +371,14 @@ def test_tiff_predictor(): @pytest.mark.enable_socket() def test_rgba(): """Decode rgb with transparency""" - reader = PdfReader(BytesIO(get_data_from_url(name="tika-972174.pdf"))) - data = reader.pages[0].images[0] - assert ".jp2" in data.name - similarity = image_similarity( - data.image, BytesIO(get_data_from_url(name="tika-972174_p0-im0.png")) - ) - assert similarity > 0.99 + with PILContext(): + reader = PdfReader(BytesIO(get_data_from_url(name="tika-972174.pdf"))) + data = reader.pages[0].images[0] + assert ".jp2" in data.name + similarity = image_similarity( + data.image, BytesIO(get_data_from_url(name="tika-972174_p0-im0.png")) + ) + assert similarity > 0.99 @pytest.mark.enable_socket() diff --git a/tests/test_workflows.py b/tests/test_workflows.py index cc8d5f94e..94e380dca 100644 --- a/tests/test_workflows.py +++ b/tests/test_workflows.py @@ -26,7 +26,7 @@ read_object, ) -from . import get_data_from_url, normalize_warnings +from . import PILContext, get_data_from_url, normalize_warnings TESTS_ROOT = Path(__file__).parent.resolve() PROJECT_ROOT = TESTS_ROOT.parent @@ -672,12 +672,13 @@ def test_image_extraction(url, name): if not root.exists(): root.mkdir() - for page in reader.pages: - for image in page.images: - filename = root / image.name - with open(filename, "wb") as img: - img.write(image.data) - images_extracted.append(filename) + with PILContext(): + for page in reader.pages: + for image in page.images: + filename = root / image.name + with open(filename, "wb") as img: + img.write(image.data) + images_extracted.append(filename) # Cleanup do_cleanup = True # set this to False for manual inspection