product-54354138-1/tools/images.py

"""Stdlib-only inspection and construction of the image formats we accept.

The CLI decides what to keep from the **downloaded bytes** and never from the
metadata a source reports about a file: a search hit that claims to be a
``image/png`` of 800x600 pixels is kept only when the bytes really are a PNG of
at least the requested size.  This module is the single place that knows how to
read the real format and the real pixel dimensions out of those bytes.

It deliberately depends on nothing but the standard library:

* :func:`detect_format` recognises the container signatures of the two accepted
  formats (``png``, ``ico``) plus the formats a real image source commonly
  serves instead (``jpeg``, ``gif``, ``bmp``, ``webp``, ``svg``, ``tiff``,
  ``pdf``).  Every recognised-but-different format, and every payload that is
  not a recognised image at all, counts as *another format* rather than an
  access problem: "the source has images, but none of them are the format you
  asked for" is a different answer from "the source could not be used".
* :func:`read_png_dimensions` reads the ``IHDR`` header of a PNG.
* :func:`read_ico_dimensions` reads the ICONDIR/ICONDIRENTRY directory of an
  ICO file (the directory stores the real sizes, including the ``0``-means-256
  encoding).
* :func:`inspect_image` combines both into an :class:`ImageInfo`.
* :func:`png_bytes` and :func:`ico_bytes` *build* the same formats so tests (and
  callers) can produce real payloads without Pillow or any other dependency.

:class:`ImageFormatError` -- a :class:`~tools.errors.ToolkitError` subclass with
a stable ``code`` -- is reserved for bytes that are genuinely unusable: an
unknown container, or a ``png``/``ico`` payload whose header is truncated or
corrupt.  A readable payload of another format is never an error, so the CLI
keeps a working search out of the ``access_error`` status.
"""

from __future__ import annotations

import struct
import zlib
from dataclasses import dataclass
from typing import Any, Iterable, Sequence

from .errors import ToolkitError

__all__ = [
    "BMP",
    "GIF",
    "ICO",
    "ICO_ENTRY_SIZE",
    "ICO_HEADER_SIZE",
    "ICO_MAX_DIMENSION",
    "ImageFormatError",
    "ImageInfo",
    "JPEG",
    "KNOWN_FORMATS",
    "PDF",
    "PNG",
    "PNG_SIGNATURE",
    "SVG",
    "TIFF",
    "WEBP",
    "detect_format",
    "ico_bytes",
    "ico_entries",
    "inspect_image",
    "png_bytes",
    "read_ico_dimensions",
    "read_png_dimensions",
]

#: The two formats the CLI accepts as search requests.
PNG = "png"
ICO = "ico"
#: Formats that are recognised (so they can be reported as "another format")
#: but never accepted by the CLI.
JPEG = "jpeg"
GIF = "gif"
BMP = "bmp"
WEBP = "webp"
#: Vector and document formats a live image source serves instead of a bitmap.
#: Naming them lets a search say *which* other formats the topic offered instead
#: of degrading into an access problem.
SVG = "svg"
TIFF = "tiff"
PDF = "pdf"

#: Every format :func:`detect_format` can name.
KNOWN_FORMATS: tuple[str, ...] = (PNG, ICO, JPEG, GIF, BMP, WEBP, SVG, TIFF, PDF)

#: The 8-byte signature every PNG file starts with.
PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n"

#: ICONDIR: reserved (0), type (1 = icon) and the entry count.
ICO_HEADER_SIZE = 6
#: One ICONDIRENTRY: 16 bytes of width/height/colours/planes/bpp/size/offset.
ICO_ENTRY_SIZE = 16
#: A single ICO dimension byte of 0 encodes 256 pixels.
ICO_MAX_DIMENSION = 256

#: Magic prefixes of the recognised but unsupported formats.
_JPEG_SIGNATURE = b"\xff\xd8\xff"
_GIF_SIGNATURES = (b"GIF87a", b"GIF89a")
_BMP_SIGNATURE = b"BM"
_RIFF_SIGNATURE = b"RIFF"
_WEBP_SIGNATURE = b"WEBP"
_TIFF_SIGNATURES = (b"II*\x00", b"MM\x00*")
_PDF_SIGNATURE = b"%PDF-"
#: How many leading bytes are searched for an ``<svg`` root element.
_SVG_PROBE = 1024

#: Number of payload bytes worth quoting in an error detail.
_SIGNATURE_PROBE = 16


class ImageFormatError(ToolkitError):
    """Raised when downloaded bytes are not a readable image of a known format.

    The bytes could not be attributed to any format the toolkit understands, or
    they claim to be an accepted format but their header is truncated/corrupt.
    It derives from :class:`~tools.errors.ToolkitError`, so the CLI reports it
    as an access problem (the source handed us something unusable) rather than
    as a "no such format" / "no such size" verdict about the topic.
    """

    code = "image_unreadable"

    def __init__(self, message: str, *, details: dict[str, Any] | None = None) -> None:
        super().__init__(message, details=details)


@dataclass(frozen=True)
class ImageInfo:
    """The real format and real pixel size of one payload."""

    format: str
    width: int
    height: int

    def to_dict(self) -> dict[str, Any]:
        """JSON-ready projection used in CLI items."""

        return {"format": self.format, "width": self.width, "height": self.height}

    @property
    def area(self) -> int:
        return self.width * self.height


def _probe(payload: bytes) -> bytes:
    return bytes(payload[:_SIGNATURE_PROBE])


def _looks_like_svg(data: bytes) -> bool:
    """True when *data* starts with an SVG root element (or its XML prologue)."""

    head = data[:_SVG_PROBE]
    if head.startswith(b"\xef\xbb\xbf"):
        head = head[3:]
    head = head.lstrip()
    if head.lower().startswith(b"<svg"):
        return True
    return head.startswith(b"<?xml") and b"<svg" in data[:_SVG_PROBE].lower()


def detect_format(payload: bytes) -> str | None:
    """Return the real container format of *payload*, or ``None``.

    A recognised format other than ``png``/``ico`` is a perfectly good answer:
    the source does have images, just not in the requested format.  ``None``
    means the bytes are not a recognised image at all -- still "not the requested
    format" from the caller's point of view, and never an access problem on its
    own.
    """

    if not isinstance(payload, (bytes, bytearray, memoryview)):
        return None
    data = bytes(payload)
    if data.startswith(PNG_SIGNATURE):
        return PNG
    if data[:4] == b"\x00\x00\x01\x00":
        return ICO
    if data.startswith(_JPEG_SIGNATURE):
        return JPEG
    if data.startswith(_GIF_SIGNATURES):
        return GIF
    if data.startswith(_BMP_SIGNATURE) and len(data) >= 14:
        return BMP
    if data.startswith(_RIFF_SIGNATURE) and data[8:12] == _WEBP_SIGNATURE:
        return WEBP
    if data.startswith(_TIFF_SIGNATURES):
        return TIFF
    if data.startswith(_PDF_SIGNATURE):
        return PDF
    if _looks_like_svg(data):
        return SVG
    return None


def read_png_dimensions(payload: bytes) -> tuple[int, int]:
    """Return ``(width, height)`` from the ``IHDR`` chunk of a PNG payload."""

    data = bytes(payload)
    if not data.startswith(PNG_SIGNATURE):
        raise ImageFormatError(
            "payload does not start with the PNG signature",
            details={"signature": _probe(data).hex()},
        )
    offset = len(PNG_SIGNATURE)
    if len(data) < offset + 8:
        raise ImageFormatError("PNG header is truncated")
    (length,) = struct.unpack(">I", data[offset : offset + 4])
    chunk_type = data[offset + 4 : offset + 8]
    if chunk_type != b"IHDR":
        raise ImageFormatError(
            "PNG payload does not start with an IHDR chunk",
            details={"first_chunk": chunk_type.decode("latin-1")},
        )
    if length != 13:
        raise ImageFormatError(
            "PNG IHDR chunk has an unexpected length",
            details={"length": length},
        )
    header = data[offset + 8 : offset + 8 + length]
    if len(header) != length:
        raise ImageFormatError("PNG IHDR chunk is truncated")
    width, height = struct.unpack(">II", header[:8])
    if width < 1 or height < 1:
        raise ImageFormatError(
            "PNG image declares no pixels",
            details={"width": width, "height": height},
        )
    return int(width), int(height)


def ico_entries(payload: bytes) -> tuple[tuple[int, int, int, int], ...]:
    """Return ``(width, height, size, offset)`` for every ICO directory entry."""

    data = bytes(payload)
    if len(data) < ICO_HEADER_SIZE:
        raise ImageFormatError("ICO header is truncated")
    reserved, kind, count = struct.unpack("<HHH", data[:ICO_HEADER_SIZE])
    if reserved != 0 or kind != 1 or count < 1:
        raise ImageFormatError(
            "payload is not an ICO image",
            details={"reserved": reserved, "type": kind, "count": count},
        )
    directory_end = ICO_HEADER_SIZE + count * ICO_ENTRY_SIZE
    if len(data) < directory_end:
        raise ImageFormatError(
            "ICO directory is truncated",
            details={"count": count, "bytes": len(data)},
        )
    entries: list[tuple[int, int, int, int]] = []
    for index in range(count):
        start = ICO_HEADER_SIZE + index * ICO_ENTRY_SIZE
        raw_width, raw_height, _colours, _reserved, _planes, _bits, size, offset = struct.unpack(
            "<BBBBHHII", data[start : start + ICO_ENTRY_SIZE]
        )
        width = raw_width or ICO_MAX_DIMENSION
        height = raw_height or ICO_MAX_DIMENSION
        if size < 1 or offset < directory_end or offset + size > len(data):
            raise ImageFormatError(
                "ICO directory entry points outside the payload",
                details={"index": index, "size": size, "offset": offset, "bytes": len(data)},
            )
        entries.append((width, height, size, offset))
    return tuple(entries)


def read_ico_dimensions(payload: bytes) -> tuple[int, int]:
    """Return the dimensions of the largest image stored in an ICO payload.

    The ICONDIRENTRY directory is authoritative: it carries one width/height pair
    per stored image, so no decoding of the embedded bitmap or PNG is needed.
    """

    entries = ico_entries(payload)
    width, height, _size, _offset = max(entries, key=lambda entry: entry[0] * entry[1])
    return int(width), int(height)


def inspect_image(payload: bytes) -> ImageInfo:
    """Return the real format and pixel size of *payload*.

    Raises :class:`ImageFormatError` when the bytes are not an image at all, or
    when they are not one of the two formats whose dimensions can be read.
    """

    image_format = detect_format(payload)
    if image_format is None:
        raise ImageFormatError(
            "downloaded payload is not a recognised image format",
            details={"signature": _probe(bytes(payload)).hex()},
        )
    if image_format == PNG:
        width, height = read_png_dimensions(payload)
        return ImageInfo(format=PNG, width=width, height=height)
    if image_format == ICO:
        width, height = read_ico_dimensions(payload)
        return ImageInfo(format=ICO, width=width, height=height)
    raise ImageFormatError(
        f"pixel dimensions are only readable for png and ico, not {image_format}",
        details={"format": image_format},
    )


# ---------------------------------------------------------------------------
# Builders: the same formats, produced with the standard library only
# ---------------------------------------------------------------------------


def _png_chunk(chunk_type: bytes, chunk_data: bytes) -> bytes:
    crc = zlib.crc32(chunk_type + chunk_data) & 0xFFFFFFFF
    return struct.pack(">I", len(chunk_data)) + chunk_type + chunk_data + struct.pack(">I", crc)


def png_bytes(
    width: int,
    height: int,
    *,
    rgba: tuple[int, int, int, int] = (255, 255, 255, 255),
    rows: Sequence[bytes] | None = None,
    level: int = 6,
) -> bytes:
    """Build a valid 8-bit RGBA PNG of exactly ``width`` x ``height`` pixels.

    ``rows`` may supply raw (filter-0) RGBA scanlines; otherwise every pixel is
    filled with ``rgba``.  The result is a real PNG file: a correct signature,
    a 13-byte ``IHDR``, a zlib-compressed ``IDAT`` and an empty ``IEND``.
    """

    if isinstance(width, bool) or isinstance(height, bool):
        raise ValueError("width and height must be integers")
    if width < 1 or height < 1:
        raise ValueError("width and height must be >= 1")
    if rows is None:
        pixel = bytes(rgba)
        line = pixel * width
        raw = bytearray()
        for _ in range(height):
            raw.append(0)
            raw += line
        raw_bytes = bytes(raw)
    else:
        if len(rows) != height:
            raise ValueError("rows must contain exactly height scanlines")
        raw = bytearray()
        for row in rows:
            if len(row) != width * 4:
                raise ValueError("each scanline must hold width * 4 bytes")
            raw.append(0)
            raw += row
        raw_bytes = bytes(raw)
    header = struct.pack(">IIBBBBB", width, height, 8, 6, 0, 0, 0)
    return (
        PNG_SIGNATURE
        + _png_chunk(b"IHDR", header)
        + _png_chunk(b"IDAT", zlib.compress(raw_bytes, level))
        + _png_chunk(b"IEND", b"")
    )


def ico_bytes(
    width: int,
    height: int,
    *,
    rgba: tuple[int, int, int, int] = (255, 255, 255, 255),
    frames: Iterable[tuple[int, int]] | None = None,
) -> bytes:
    """Build an ICO payload whose directory lists the requested sizes.

    By default the ICO stores one embedded PNG of ``width`` x ``height`` pixels
    and the directory entry repeats those dimensions (``0`` encodes 256).  Pass
    ``frames`` to describe several directory entries instead, for example
    ``ico_bytes(48, 48, frames=[(16, 16), (48, 48)])``: the dimensions of the
    *largest* entry then become the reported size, which is exactly what
    :func:`read_ico_dimensions` returns.
    """

    sizes = list(frames) if frames is not None else [(width, height)]
    if not sizes:
        raise ValueError("an ICO payload needs at least one frame")
    payloads: list[bytes] = []
    for frame_width, frame_height in sizes:
        if not 1 <= frame_width <= ICO_MAX_DIMENSION:
            raise ValueError("ICO widths must be between 1 and 256 pixels")
        if not 1 <= frame_height <= ICO_MAX_DIMENSION:
            raise ValueError("ICO heights must be between 1 and 256 pixels")
        payloads.append(png_bytes(frame_width, frame_height, rgba=rgba))
    header = struct.pack("<HHH", 0, 1, len(payloads))
    offset = ICO_HEADER_SIZE + ICO_ENTRY_SIZE * len(payloads)
    entries = bytearray()
    for (frame_width, frame_height), payload in zip(sizes, payloads):
        entries += struct.pack(
            "<BBBBHHII",
            frame_width % ICO_MAX_DIMENSION,
            frame_height % ICO_MAX_DIMENSION,
            0,
            0,
            1,
            32,
            len(payload),
            offset,
        )
        offset += len(payload)
    return header + bytes(entries) + b"".join(payloads)