diff --git a/docling/.agents/skills/docling/references/python-sdk.md b/docling/.agents/skills/docling/references/python-sdk.md index be17238dad..e44ae5fd95 100644 --- a/docling/.agents/skills/docling/references/python-sdk.md +++ b/docling/.agents/skills/docling/references/python-sdk.md @@ -465,7 +465,8 @@ Or save with `ImageRefMode.EMBEDDED` so the JSON carries its images. HTML image downloads are off by default. Enable them on the HTML backend options; downloads only connect to public addresses, re-check every redirect -and stop at `max_remote_image_bytes`. +and stop at `max_remote_image_bytes`. An image larger than `max_image_pixels` +(width x height, default 89,478,485) is skipped with a warning. ```python from docling.datamodel.backend_options import HTMLBackendOptions diff --git a/docling/backend/html_backend.py b/docling/backend/html_backend.py index da645055bc..9794a8718f 100644 --- a/docling/backend/html_backend.py +++ b/docling/backend/html_backend.py @@ -79,6 +79,7 @@ detect_code_language, normalize_code_language, ) +from docling.utils.image_utils import check_image_pixels _BS4_AVAILABLE: bool = False _BS4_IMPORT_ERROR: ImportError | None = None @@ -536,6 +537,7 @@ def __init__( enable_remote_fetch=options.enable_remote_fetch, max_image_data_base64_bytes=options.max_image_data_base64_bytes, max_remote_image_bytes=options.max_remote_image_bytes, + max_image_pixels=options.max_image_pixels, max_redirects=options.max_redirects, headers=options.headers, header_origins=self._get_header_origins(options, configured_base_path), @@ -5375,11 +5377,18 @@ def _create_image_ref(self, src_url: str) -> Optional[ImageRef]: return None try: image = Image.open(BytesIO(image_data)) + check_image_pixels(image.size, self.options.max_image_pixels) image.load() return ImageRef.from_pil( image, dpi=int(image.info.get("dpi", (72,))[0]) ) - except (UnidentifiedImageError, OSError, TypeError, ValueError) as exc: + except ( + Image.DecompressionBombError, + UnidentifiedImageError, + OSError, + TypeError, + ValueError, + ) as exc: warnings.warn( f"Could not process an embedded MHTML image from {src_url}: " f"{exc}" diff --git a/docling/backend/image_backend.py b/docling/backend/image_backend.py index 90da683f93..864691613d 100644 --- a/docling/backend/image_backend.py +++ b/docling/backend/image_backend.py @@ -3,6 +3,7 @@ import logging import math +import threading from collections.abc import Iterable from io import BytesIO from pathlib import Path @@ -24,7 +25,7 @@ from docling.datamodel.base_models import InputFormat from docling.datamodel.document import InputDocument from docling.exceptions import DocumentLoadError -from docling.utils.image_utils import to_rgb +from docling.utils.image_utils import check_image_pixels, to_rgb _log = logging.getLogger(__name__) @@ -200,8 +201,10 @@ class ImageDocumentBackend(PdfDocumentBackend): - Subclasses PdfDocumentBackend to satisfy pipeline type checks. - Intentionally avoids calling PdfDocumentBackend.__init__ to skip the image→PDF conversion and any pypdfium2 usage. - - Handles multi-page TIFF by extracting frames eagerly to separate - Image objects to keep thread-safety when pages process in parallel. + - Reads only the frame count when opened, so the page limits are + checked before any pixel data is decoded. Each page decodes its + frame on ``load_page`` into a separate Image object, which keeps + pages thread-safe when they process in parallel. """ def __init__( @@ -221,47 +224,52 @@ def __init__( f"Incompatible file format {self.input_format} was passed to ImageDocumentBackend." ) - # Load frames eagerly for thread-safety across pages - self._frames: list[Image.Image] = [] - self._frame_dpi: list[tuple[float, float]] = [] + # The source stays open until unload(); the lock serializes the seeks + # on it, because pages can load from several threads. + self._lock = threading.Lock() + self._source: Image.Image | None = None + self._frame_count = 0 try: - with Image.open(self.path_or_stream) as img: # type: ignore[arg-type] - # Handle multi-frame and single-frame images - # - multiframe formats: TIFF, GIF, ICO - # - singleframe formats: JPEG (.jpg, .jpeg), PNG (.png), BMP, WEBP (unless animated), HEIC - frame_count = getattr(img, "n_frames", 1) - - if frame_count > 1: - for i in range(frame_count): - img.seek(i) - frame, dpi = _oriented_frame(img) - self._frame_dpi.append(dpi) - self._frames.append(to_rgb(frame)) - else: - frame, dpi = _oriented_frame(img) - self._frame_dpi.append(dpi) - self._frames.append(to_rgb(frame)) + self._source = Image.open(self.path_or_stream) # type: ignore[arg-type] + # Multi-frame formats: TIFF, GIF, ICO, animated WEBP/PNG. Other + # formats have no n_frames attribute and hold one frame. + self._frame_count = getattr(self._source, "n_frames", 1) + # The first frame is the current one after open (and after + # n_frames, which seeks back): reject a bad first frame here, + # without decoding it. + check_image_pixels(self._source.size, self.options.max_image_pixels) + _get_frame_dpi(self._source) except Exception as e: - for frame in self._frames: - frame.close() - self._frames = [] - self._frame_dpi = [] + self._close_source() raise DocumentLoadError( - f"Could not load image for document {self.file}" + f"Could not load image for document {self.file}: {e}" ) from e + def _close_source(self) -> None: + if self._source is not None: + self._source.close() + self._source = None + self._frame_count = 0 + def is_valid(self) -> bool: - return len(self._frames) > 0 + return self._frame_count > 0 def page_count(self) -> int: - return len(self._frames) + return self._frame_count def load_page(self, page_no: int) -> _ImagePageBackend: - if not (0 <= page_no < len(self._frames)): + if not (0 <= page_no < self._frame_count): raise IndexError(f"Page index out of range: {page_no}") - return _ImagePageBackend( - self._frames[page_no], page_no, self._frame_dpi[page_no] - ) + with self._lock: + assert self._source is not None + if self._source.tell() != page_no: + self._source.seek(page_no) + # Frames of a multi-page TIFF have their own size, so each one is + # checked before it is decoded. + check_image_pixels(self._source.size, self.options.max_image_pixels) + frame, dpi = _oriented_frame(self._source) + image = to_rgb(frame) + return _ImagePageBackend(image, page_no, dpi) @classmethod def supported_formats(cls) -> set[InputFormat]: @@ -273,8 +281,6 @@ def supports_pagination(cls) -> bool: return True def unload(self): - for frame in self._frames: - frame.close() - self._frames = [] - self._frame_dpi = [] + with self._lock: + self._close_source() super().unload() diff --git a/docling/backend/md_backend.py b/docling/backend/md_backend.py index 99cbcaf23a..7e77cb8054 100644 --- a/docling/backend/md_backend.py +++ b/docling/backend/md_backend.py @@ -1112,6 +1112,7 @@ def _get_image_loader(self) -> ImageResourceLoader: enable_local_fetch=md_options.enable_local_fetch, enable_remote_fetch=md_options.enable_remote_fetch, max_image_data_base64_bytes=md_options.max_image_data_base64_bytes, + max_image_pixels=md_options.max_image_pixels, ) return self._image_loader diff --git a/docling/backend/utils/image_resource_loader.py b/docling/backend/utils/image_resource_loader.py index 5d498761a9..e1e1a99715 100644 --- a/docling/backend/utils/image_resource_loader.py +++ b/docling/backend/utils/image_resource_loader.py @@ -5,8 +5,8 @@ Turns an image source (a URI, a local file, or a remote URL) into a DoclingDocument `ImageRef`, enforcing the relevant safety limits: a path-traversal guard when resolving -relative paths, the `enable_local_fetch` / `enable_remote_fetch` toggles, and the -base64 and remote download size caps. +relative paths, the `enable_local_fetch` / `enable_remote_fetch` toggles, the +base64 and remote download size caps, and a limit on the decoded image pixels. """ import base64 @@ -32,7 +32,9 @@ from pydantic import ValidationError from requests.utils import get_environ_proxies, select_proxy +from docling.datamodel.backend_options import DEFAULT_MAX_IMAGE_PIXELS from docling.exceptions import OperationNotAllowed +from docling.utils.image_utils import check_image_pixels _log = logging.getLogger(__name__) @@ -267,6 +269,7 @@ def __init__( enable_remote_fetch: bool = False, max_image_data_base64_bytes: int = 20 * 1024 * 1024, max_remote_image_bytes: int = 20 * 1024 * 1024, + max_image_pixels: int = DEFAULT_MAX_IMAGE_PIXELS, max_redirects: int = 5, headers: Optional[dict[str, str]] = None, header_origins: Iterable[str] = (), @@ -275,6 +278,7 @@ def __init__( self.enable_remote_fetch = enable_remote_fetch self.max_image_data_base64_bytes = max_image_data_base64_bytes self.max_remote_image_bytes = max_remote_image_bytes + self.max_image_pixels = max_image_pixels self.max_redirects = max_redirects self.headers = headers # Only requests to these origins carry the configured `headers`. @@ -348,8 +352,12 @@ def create_image_ref( img_data = self.load_image_data(src_url, base_path) if img_data: img = Image.open(BytesIO(img_data)) + # Image.open reads only the header: check the size before + # ImageRef.from_pil decodes the pixels. + check_image_pixels(img.size, self.max_image_pixels) return ImageRef.from_pil(img, dpi=int(img.info.get("dpi", (72,))[0])) except ( + Image.DecompressionBombError, requests.RequestException, urllib3.exceptions.HTTPError, ValidationError, @@ -454,6 +462,14 @@ def load_image_data( return self.fetch_remote(src_loc).content elif src_loc.startswith("data:"): encoded_data = re.sub(r"^data:image/.+;base64,", "", src_loc) + # Check the size before decoding: 4 base64 characters hold 3 + # bytes, less one byte for each "=" padding character. + padding = len(encoded_data[-2:]) - len(encoded_data[-2:].rstrip("=")) + decoded_size = len(encoded_data) * 3 // 4 - padding + if decoded_size > self.max_image_data_base64_bytes: + raise ValueError( + f"Decoded image exceeds size limit of {self.max_image_data_base64_bytes} bytes." + ) decoded_data = base64.b64decode(encoded_data) if len(decoded_data) > self.max_image_data_base64_bytes: diff --git a/docling/datamodel/backend_options.py b/docling/datamodel/backend_options.py index 0e1a84e73a..adeb33789b 100644 --- a/docling/datamodel/backend_options.py +++ b/docling/datamodel/backend_options.py @@ -20,6 +20,11 @@ model_validator, ) +# Default limit on the pixels of one decoded image or image frame. It is the +# default of Pillow's ``Image.MAX_IMAGE_PIXELS`` (about 268 MB once decoded to +# RGB), fixed here so that the limit does not depend on global Pillow state. +DEFAULT_MAX_IMAGE_PIXELS = 89_478_485 + class BaseBackendOptions(BaseModel): """Common options for all declarative document backends.""" @@ -196,6 +201,13 @@ class HTMLBackendOptions(BaseBackendOptions): 20 * 1024 * 1024, # 20 MB description="The maximum number of bytes for remote image downloads.", ) + max_image_pixels: PositiveInt = Field( + DEFAULT_MAX_IMAGE_PIXELS, + description=( + "The maximum number of pixels (width x height) of an embedded or " + "fetched image. Larger images are skipped with a warning." + ), + ) max_redirects: Annotated[int, Field(ge=0)] = Field( 5, description="Maximum number of HTTP redirects to follow when fetching remote resources. Set to 0 to disable redirects.", @@ -236,6 +248,13 @@ class MarkdownBackendOptions(TextBackendOptions): 20 * 1024 * 1024, # 20 MB description="The maximum number of base64 data bytes that the backend will accept.", ) + max_image_pixels: PositiveInt = Field( + DEFAULT_MAX_IMAGE_PIXELS, + description=( + "The maximum number of pixels (width x height) of an embedded or " + "fetched image. Larger images are skipped with a warning." + ), + ) class JatsBackendOptions(BaseBackendOptions): @@ -306,6 +325,14 @@ class PdfBackendOptions(BaseBackendOptions): "decoding costs time and memory, so it is off by default." ), ) + max_image_pixels: PositiveInt = Field( + DEFAULT_MAX_IMAGE_PIXELS, + description=( + "The maximum number of pixels (width x height) of one page of an " + "image input, such as one frame of a multi-page TIFF. Used by the " + "image backend; an image with a larger page is rejected." + ), + ) class ThreadedDoclingParseBackendOptions(PdfBackendOptions): diff --git a/docling/utils/image_utils.py b/docling/utils/image_utils.py index a952422552..06ece1058f 100644 --- a/docling/utils/image_utils.py +++ b/docling/utils/image_utils.py @@ -4,6 +4,20 @@ from PIL import Image +def check_image_pixels(size: tuple[int, int], max_pixels: int) -> None: + """Raise ``ValueError`` if an image of ``size`` exceeds ``max_pixels``. + + Pillow reads the dimensions from the image header, so call this after + ``Image.open`` (or ``seek``) and before the pixel data is decoded. + """ + width, height = size + if width * height > max_pixels: + raise ValueError( + f"Image of {width}x{height} pixels exceeds the limit of " + f"{max_pixels} pixels." + ) + + def to_rgb(image: Image.Image) -> Image.Image: """Convert an image to RGB, flattening any transparency onto white. diff --git a/docs/usage/advanced_options.md b/docs/usage/advanced_options.md index b548926d63..5275dcf2a3 100644 --- a/docs/usage/advanced_options.md +++ b/docs/usage/advanced_options.md @@ -568,7 +568,9 @@ The HTML backend only downloads images referenced by a page when you opt in with `--html-image-fetch remote`). Downloads connect only to public, globally routable addresses: every address of a host is checked, every redirect is checked again before it is followed (up to `max_redirects`), and downloads stop -at `max_remote_image_bytes`. When `render_page=True`, the browser requests +at `max_remote_image_bytes`. An image larger than `max_image_pixels` (width x +height, default 89,478,485) is skipped with a warning, and the rest of the +document converts; the Markdown backend has the same option. When `render_page=True`, the browser requests remote resources through the same download path, and navigating the page away from the source document is refused. @@ -620,6 +622,27 @@ converter = DocumentConverter() result = converter.convert(source, max_num_pages=100, max_file_size=20971520) ``` +For image inputs, the page count is checked against `max_num_pages` before any +page is decoded, and each page (for example each frame of a multi-page TIFF) is +decoded only when it is processed. The `max_image_pixels` backend option limits +the pixels (width x height) of one page; the default, 89,478,485, is the default +of Pillow's `Image.MAX_IMAGE_PIXELS`. A larger first page rejects the document, +and a larger later page fails that page: + +```python +from docling.datamodel.backend_options import PdfBackendOptions +from docling.datamodel.base_models import InputFormat +from docling.document_converter import DocumentConverter, ImageFormatOption + +converter = DocumentConverter( + format_options={ + InputFormat.IMAGE: ImageFormatOption( + backend_options=PdfBackendOptions(max_image_pixels=50_000_000) + ) + } +) +``` + ## Convert from binary PDF streams You can convert PDFs from a binary stream instead of from the filesystem as follows: diff --git a/tests/test_backend_html.py b/tests/test_backend_html.py index a1818b2cb5..f991ef316a 100644 --- a/tests/test_backend_html.py +++ b/tests/test_backend_html.py @@ -20,6 +20,7 @@ RichTableCell, ) from docling_core.types.doc.document import ContentLayer +from PIL import Image from pydantic import AnyUrl, ValidationError from docling.backend.html_backend import ( @@ -31,7 +32,7 @@ validate_url_safety as _validate_url_safety, ) from docling.datamodel.backend_options import HTMLBackendOptions -from docling.datamodel.base_models import ConversionStatus, InputFormat +from docling.datamodel.base_models import ConversionStatus, DocumentStream, InputFormat from docling.datamodel.document import ( ConversionResult, DoclingDocument, @@ -2010,6 +2011,33 @@ def test_load_image_data_enforces_data_uri_size_limit(): backend._load_image_data(data_uri) +def test_image_over_pixel_limit_is_skipped(): + def png_data_uri(width: int, height: int) -> str: + buffer = BytesIO() + Image.new("RGB", (width, height), (255, 0, 0)).save(buffer, format="PNG") + return "data:image/png;base64," + base64.b64encode(buffer.getvalue()).decode() + + html = ( + "

Before.

" + f'' + "

After.

" + ) + converter = _create_html_converter( + HTMLBackendOptions(fetch_images=True, max_image_pixels=63) + ) + with pytest.warns(UserWarning, match="exceeds the limit of 63 pixels"): + result = converter.convert( + DocumentStream(name="test.html", stream=BytesIO(html.encode())) + ) + doc = result.document + + assert result.status == ConversionStatus.SUCCESS + assert len(doc.pictures) == 2 + assert doc.pictures[0].image is None + assert doc.pictures[1].get_image(doc).size == (7, 5) + assert [text.text for text in doc.texts] == ["Before.", "After."] + + def test_anchor_fragment_links_with_source_uri(): """Fragment-only hrefs must not be mangled when source_uri is set.""" html_path = Path("tests/data/html/sources/hyperlink_06.html") diff --git a/tests/test_backend_image_native.py b/tests/test_backend_image_native.py index 664d84246c..a8a6c230dd 100644 --- a/tests/test_backend_image_native.py +++ b/tests/test_backend_image_native.py @@ -2,13 +2,14 @@ # SPDX-License-Identifier: MIT from io import BytesIO -from unittest.mock import MagicMock import pytest from docling_core.types.doc import BoundingBox, CoordOrigin from PIL import Image, TiffImagePlugin +from docling.backend import image_backend from docling.backend.image_backend import ImageDocumentBackend, _ImagePageBackend +from docling.datamodel.backend_options import PdfBackendOptions from docling.datamodel.base_models import DocumentStream, InputFormat from docling.datamodel.document import ( InputDocument, @@ -16,6 +17,7 @@ _DummyBackend, get_input_rejection_cause, ) +from docling.datamodel.settings import DocumentLimits from docling.document_converter import DocumentConverter, ImageFormatOption from docling.document_extractor import DocumentExtractor from docling.exceptions import DocumentLoadError @@ -419,36 +421,17 @@ def test_invalid_explicit_dpi_rejects_document(): assert isinstance(get_input_rejection_cause(in_doc), DocumentLoadError) -def test_source_image_is_closed_after_backend_init(tmp_path, monkeypatch): +def test_source_image_is_closed_on_unload(tmp_path, monkeypatch): image_path = tmp_path / "test.png" Image.new("RGB", (32, 32), (10, 20, 30)).save(image_path) opened_images = [] original_open = Image.open - class TrackingImage: - def __init__(self, image): - self._image = image - self.closed = False - - def __getattr__(self, attr): - return getattr(self._image, attr) - - def close(self): - self.closed = True - return self._image.close() - - def __enter__(self): - return self - - def __exit__(self, exc_type, exc, tb): - self.close() - return False - def tracking_open(*args, **kwargs): - tracked_image = TrackingImage(original_open(*args, **kwargs)) - opened_images.append(tracked_image) - return tracked_image + image = original_open(*args, **kwargs) + opened_images.append(image) + return image input_doc = InputDocument( path_or_stream=image_path, @@ -462,25 +445,72 @@ def tracking_open(*args, **kwargs): in_doc=input_doc, path_or_stream=image_path, ) + page_image = backend.load_page(0).get_page_image() + backend.unload() assert len(opened_images) == 1 - assert opened_images[0].closed is True - backend.unload() + assert opened_images[0].fp is None + # The decoded page stays usable after the source is closed. + assert page_image.getpixel((0, 0)) == (10, 20, 30) -def test_unload_closes_cached_frames(): - stream = _make_multipage_tiff_stream(num_pages=3, size=(32, 32)) - doc_backend = _get_backend_from_stream(stream) +def test_frame_count_over_max_num_pages_is_rejected_before_decoding(monkeypatch): + decoded = [] + original_to_rgb = image_backend.to_rgb - tracked_closers = [] - for frame in doc_backend._frames: - closer = MagicMock(wraps=frame.close) - frame.close = closer - tracked_closers.append(closer) + def counting_to_rgb(image): + decoded.append(image.size) + return original_to_rgb(image) + + monkeypatch.setattr(image_backend, "to_rgb", counting_to_rgb) + stream = _make_multipage_tiff_stream(num_pages=5, size=(16, 16)) + in_doc = InputDocument( + path_or_stream=stream.stream, + format=InputFormat.IMAGE, + backend=ImageDocumentBackend, + filename=stream.name, + limits=DocumentLimits(max_num_pages=2), + ) + + assert in_doc.valid is False + assert in_doc.page_count == 5 + assert "max_num_pages limit of 2" in in_doc._rejection.message + assert decoded == [] + in_doc._backend.unload() - doc_backend.unload() - assert doc_backend._frames == [] - assert doc_backend._frame_dpi == [] - for closer in tracked_closers: - closer.assert_called_once() +def test_image_over_pixel_limit_is_rejected(): + stream = _make_png_stream(width=64, height=48) + in_doc = InputDocument( + path_or_stream=stream.stream, + format=InputFormat.IMAGE, + backend=ImageDocumentBackend, + backend_options=PdfBackendOptions(max_image_pixels=64 * 48 - 1), + filename=stream.name, + ) + + assert in_doc.valid is False + cause = get_input_rejection_cause(in_doc) + assert isinstance(cause, DocumentLoadError) + assert "exceeds the limit of 3071 pixels" in str(cause) + + +def test_tiff_frame_over_pixel_limit_is_rejected_on_load(): + frames = [Image.new("RGB", (16, 16)), Image.new("RGB", (64, 64))] + buf = BytesIO() + frames[0].save(buf, format="TIFF", save_all=True, append_images=frames[1:]) + buf.seek(0) + in_doc = InputDocument( + path_or_stream=buf, + format=InputFormat.IMAGE, + backend=ImageDocumentBackend, + backend_options=PdfBackendOptions(max_image_pixels=1000), + filename="test.tiff", + ) + doc_backend = in_doc._backend + + assert in_doc.valid is True + assert doc_backend.load_page(0).get_size().width == 16 + with pytest.raises(ValueError, match="64x64 pixels exceeds the limit"): + doc_backend.load_page(1) + doc_backend.unload() diff --git a/tests/test_backend_markdown_images.py b/tests/test_backend_markdown_images.py index 57c9c2d6bf..e843575035 100644 --- a/tests/test_backend_markdown_images.py +++ b/tests/test_backend_markdown_images.py @@ -87,6 +87,22 @@ def test_mixed_html_keeps_size_rejected_image_placeholder(html_block: str) -> No assert doc.pictures[0].caption_text(doc) == "caption" +def test_image_over_pixel_limit_is_skipped() -> None: + markdown = ( + f"Before.\n\n![big]({_png_data_uri(8, 8)})\n\n" + f"![small]({_png_data_uri(7, 5)})\n\nAfter.\n" + ) + with pytest.warns(UserWarning, match="exceeds the limit of 63 pixels"): + doc = _convert_markdown( + markdown, MarkdownBackendOptions(fetch_images=True, max_image_pixels=63) + ) + + assert len(doc.pictures) == 2 + assert doc.pictures[0].image is None + assert doc.pictures[1].get_image(doc).size == (7, 5) + assert [text.text for text in doc.texts] == ["Before.", "big", "small", "After."] + + @pytest.mark.parametrize("fetch_images", [False, True]) def test_original_html_images_keep_the_normal_loader(fetch_images: bool) -> None: markdown = ( diff --git a/tests/test_image_resource_loader.py b/tests/test_image_resource_loader.py index 37ba0d6859..04dd602b82 100644 --- a/tests/test_image_resource_loader.py +++ b/tests/test_image_resource_loader.py @@ -3,11 +3,16 @@ """Unit tests for the shared image-resource loader and its safety limits.""" +import base64 import socket +import struct import threading +import zlib +from io import BytesIO from urllib.parse import quote import pytest +from PIL import Image from docling.backend.utils.image_resource_loader import ( ImageResourceLoader, @@ -242,3 +247,49 @@ def test_load_image_data_local_requires_base_path(): loader = ImageResourceLoader(enable_local_fetch=True) with pytest.raises(OperationNotAllowed, match="requires base_path"): loader.load_image_data("/some/where/image.png", None) + + +def _png_header_data_uri(width: int, height: int) -> str: + """A PNG that declares ``width`` x ``height`` pixels but has no pixel data.""" + + def chunk(kind: bytes, data: bytes) -> bytes: + body = kind + data + return struct.pack(">I", len(data)) + body + struct.pack(">I", zlib.crc32(body)) + + png = ( + b"\x89PNG\r\n\x1a\n" + + chunk(b"IHDR", struct.pack(">IIBBBBB", width, height, 8, 0, 0, 0, 0)) + + chunk(b"IDAT", zlib.compress(b"")) + + chunk(b"IEND", b"") + ) + return "data:image/png;base64," + base64.b64encode(png).decode() + + +def _png_data_uri(width: int, height: int) -> str: + buffer = BytesIO() + Image.new("RGB", (width, height)).save(buffer, format="PNG") + return "data:image/png;base64," + base64.b64encode(buffer.getvalue()).decode() + + +def test_image_over_pixel_limit_is_skipped(): + loader = ImageResourceLoader(max_image_pixels=63) + + assert loader.create_image_ref(_png_data_uri(7, 9), None) is not None + with pytest.warns(UserWarning, match="8x8 pixels exceeds the limit of 63"): + assert loader.create_image_ref(_png_data_uri(8, 8), None) is None + + +def test_image_over_pillow_pixel_limit_is_skipped(): + # Above twice Image.MAX_IMAGE_PIXELS, Pillow refuses to open the image. + loader = ImageResourceLoader(max_image_pixels=10**9) + + with pytest.warns(UserWarning, match="400000000 pixels"): + assert loader.create_image_ref(_png_header_data_uri(20000, 20000), None) is None + + +def test_data_uri_over_size_limit_is_rejected(): + loader = ImageResourceLoader(max_image_data_base64_bytes=8) + + assert loader.load_image_data("data:image/png;base64,AAAAAAAAAAA=", None) + with pytest.raises(ValueError, match="exceeds size limit of 8 bytes"): + loader.load_image_data("data:image/png;base64,AAAAAAAAAAAAAAAA", None)