png API¶

def decode_png(source: str | os.PathLike[str] | IO[bytes] | bytes) -> Image:
    """Decode a PNG image into an :class:`Image`.

    Args:
        source: File path, file object, or raw PNG bytes.

    Returns:
        Decoded image.

    Raises:
        DecodeError: If the data is not valid PNG or uses unsupported features.
    """
    raw = _read_source(source)
    chunks = _read_chunks(raw)

    # -- IHDR -----------------------------------------------------------------
    if not chunks or chunks[0][0] != b"IHDR":
        raise DecodeError("missing IHDR chunk")
    ihdr = chunks[0][1]
    if len(ihdr) != 13:
        raise DecodeError("bad IHDR length")
    width, height = struct.unpack(">II", ihdr[:8])
    bit_depth = ihdr[8]
    color_type = ihdr[9]
    compression = ihdr[10]
    filter_method = ihdr[11]
    interlace = ihdr[12]

    if width == 0 or height == 0:
        raise DecodeError("zero-dimension image")
    if compression != 0:
        raise DecodeError(f"unsupported compression method {compression}")
    if filter_method != 0:
        raise DecodeError(f"unsupported filter method {filter_method}")
    if interlace != 0:
        raise DecodeError("interlaced PNG is not supported")
    if color_type not in _COLOR_TYPE_TO_MODE:
        raise DecodeError(f"unsupported color type {color_type}")
    if bit_depth not in (8, 16):
        raise DecodeError(f"unsupported bit depth {bit_depth}")

    mode = _COLOR_TYPE_TO_MODE[color_type]
    bpp = _bytes_per_pixel(mode, bit_depth)
    row_bytes = width * bpp

    # -- Collect IDAT and tEXt ------------------------------------------------
    idat_parts: list[bytes] = []
    metadata: dict[str, str] = {}
    for ctype, cdata in chunks:
        if ctype == b"IDAT":
            idat_parts.append(cdata)
        elif ctype == b"tEXt":
            sep = cdata.find(b"\x00")
            if sep != -1:
                key = cdata[:sep].decode("latin-1")
                val = cdata[sep + 1 :].decode("latin-1")
                metadata[key] = val

    if not idat_parts:
        raise DecodeError("no IDAT chunks found")

    # -- Decompress and unfilter ----------------------------------------------
    try:
        decompressed = zlib.decompress(b"".join(idat_parts))
    except zlib.error as exc:
        raise DecodeError(f"zlib decompression failed: {exc}") from exc

    expected_len = height * (1 + row_bytes)  # 1 filter byte per row
    if len(decompressed) != expected_len:
        raise DecodeError(
            f"decompressed size {len(decompressed)} != expected {expected_len}"
        )

    pixels = bytearray(height * row_bytes)
    prev_row: bytes | None = None
    stride = 1 + row_bytes
    for y in range(height):
        offset = y * stride
        ftype = decompressed[offset]
        row_data = decompressed[offset + 1 : offset + 1 + row_bytes]
        unfiltered = _unfilter_row(ftype, row_data, prev_row, bpp)
        pixels[y * row_bytes : (y + 1) * row_bytes] = unfiltered
        prev_row = bytes(unfiltered)

    return Image(
        width=width,
        height=height,
        data=bytes(pixels),
        mode=mode,
        bit_depth=bit_depth,
        metadata=metadata or None,
    )

def encode_png(
    image: Image,
    dest: str | os.PathLike[str] | IO[bytes] | None = None,
    *,
    filter_strategy: str = "auto",
    compression_level: int = 6,
) -> bytes | None:
    """Encode an :class:`Image` as PNG.

    Args:
        image: Image to encode.
        dest: File path, file object, or ``None`` to return bytes.
        filter_strategy: Row filter strategy — ``'auto'`` (minimum-sum
            heuristic), ``'none'``, ``'sub'``, ``'up'``, ``'avg'``, or
            ``'paeth'``.
        compression_level: zlib compression level (0–9).

    Returns:
        PNG bytes when *dest* is ``None``, otherwise ``None``.
    """
    if image.mode not in _MODE_TO_COLOR_TYPE:
        raise EncodeError(f"unsupported mode {image.mode!r}")

    color_type = _MODE_TO_COLOR_TYPE[image.mode]
    bpp = _bytes_per_pixel(image.mode, image.bit_depth)
    row_bytes = image.width * bpp

    fixed_filter: int | None = {
        "none": 0,
        "sub": 1,
        "up": 2,
        "avg": 3,
        "paeth": 4,
    }.get(filter_strategy)

    # -- Build filtered scanlines ---------------------------------------------
    total_filtered = image.height * (1 + row_bytes)
    filtered = bytearray(total_filtered)
    prev_row = b"\x00" * row_bytes
    filt_stride = 1 + row_bytes
    for y in range(image.height):
        row = image.data[y * row_bytes : (y + 1) * row_bytes]
        filt_offset = y * filt_stride
        if filter_strategy == "auto":
            ftype, fdata = _select_filter(row, prev_row, bpp)
        elif fixed_filter is not None:
            ftype = fixed_filter
            fdata = _filter_row(ftype, row, prev_row, bpp)
        else:
            raise EncodeError(f"unknown filter_strategy {filter_strategy!r}")
        filtered[filt_offset] = ftype
        filtered[filt_offset + 1 : filt_offset + 1 + row_bytes] = fdata
        prev_row = row

    compressed = zlib.compress(bytes(filtered), compression_level)

    # -- Assemble chunks ------------------------------------------------------
    ihdr_data = struct.pack(
        ">IIBBBBB",
        image.width,
        image.height,
        image.bit_depth,
        color_type,
        0,  # compression
        0,  # filter method
        0,  # interlace
    )

    parts = [_PNG_SIGNATURE, _write_chunk(b"IHDR", ihdr_data)]

    # tEXt chunks
    if image.metadata:
        for key, val in image.metadata.items():
            tdata = key.encode("latin-1") + b"\x00" + val.encode("latin-1")
            parts.append(_write_chunk(b"tEXt", tdata))

    parts.append(_write_chunk(b"IDAT", compressed))
    parts.append(_write_chunk(b"IEND", b""))

    return _write_dest(dest, b"".join(parts))

def decode_bmp(source: str | os.PathLike[str] | IO[bytes] | bytes) -> Image:
    """Decode a BMP image into an :class:`Image`.

    Args:
        source: File path, file object, or raw BMP bytes.

    Returns:
        Decoded image.

    Raises:
        DecodeError: If the data is not valid BMP or uses unsupported features.
    """
    raw = _read_source(source)

    if len(raw) < 54:
        raise DecodeError("file too small for BMP")
    if raw[:2] != _BMP_SIGNATURE:
        raise DecodeError("not a BMP file (bad signature)")

    # File header (14 bytes)
    data_offset = struct.unpack("<I", raw[10:14])[0]

    # DIB header (BITMAPINFOHEADER, 40 bytes minimum)
    dib_size = struct.unpack("<I", raw[14:18])[0]
    if dib_size < 40:
        raise DecodeError(f"unsupported DIB header size {dib_size}")

    width = struct.unpack("<i", raw[18:22])[0]
    height_raw = struct.unpack("<i", raw[22:26])[0]
    bits_per_pixel = struct.unpack("<H", raw[28:30])[0]
    compression = struct.unpack("<I", raw[30:34])[0]

    if compression != 0:
        raise DecodeError(f"unsupported BMP compression {compression}")
    if bits_per_pixel not in (24, 32):
        raise DecodeError(f"unsupported BMP bit depth {bits_per_pixel}")

    top_down = height_raw < 0
    height = abs(height_raw)

    if width <= 0 or height == 0:
        raise DecodeError("invalid BMP dimensions")

    bmp_bpp = bits_per_pixel // 8
    row_data_bytes = width * bmp_bpp
    row_stride = (row_data_bytes + 3) & ~3  # pad to 4-byte boundary

    mode = "RGBA" if bits_per_pixel == 32 else "RGB"
    out_bpp = _MODE_BPP[mode]

    pixels = bytearray(width * height * out_bpp)
    for y in range(height):
        # BMP default is bottom-up; top-down if height was negative
        src_y = y if top_down else (height - 1 - y)
        src_offset = data_offset + src_y * row_stride
        if src_offset + row_data_bytes > len(raw):
            raise DecodeError("truncated BMP pixel data")

        src_row = raw[src_offset : src_offset + row_data_bytes]
        dst_offset = y * width * out_bpp
        dst_row = bytearray(width * out_bpp)
        if bmp_bpp == 3:
            # BGR -> RGB via slice assignment (C-level ops)
            dst_row[0::3] = src_row[2::3]  # R
            dst_row[1::3] = src_row[1::3]  # G
            dst_row[2::3] = src_row[0::3]  # B
        else:
            # BGRA -> RGBA via slice assignment
            dst_row[0::4] = src_row[2::4]  # R
            dst_row[1::4] = src_row[1::4]  # G
            dst_row[2::4] = src_row[0::4]  # B
            dst_row[3::4] = src_row[3::4]  # A
        pixels[dst_offset : dst_offset + width * out_bpp] = dst_row

    return Image(width=width, height=height, data=bytes(pixels), mode=mode)

def encode_bmp(
    image: Image,
    dest: str | os.PathLike[str] | IO[bytes] | None = None,
) -> bytes | None:
    """Encode an :class:`Image` as BMP.

    Args:
        image: Image to encode (must be ``'RGB'`` or ``'RGBA'``).
        dest: File path, file object, or ``None`` to return bytes.

    Returns:
        BMP bytes when *dest* is ``None``, otherwise ``None``.
    """
    if image.mode not in ("RGB", "RGBA"):
        raise EncodeError(f"BMP encode requires RGB or RGBA mode, got {image.mode!r}")
    if image.bit_depth != 8:
        raise EncodeError("BMP encode only supports 8-bit images")

    bmp_bpp = 4 if image.mode == "RGBA" else 3
    src_bpp = _MODE_BPP[image.mode]
    bits_per_pixel = bmp_bpp * 8
    row_data_bytes = image.width * bmp_bpp
    row_stride = (row_data_bytes + 3) & ~3
    padding = row_stride - row_data_bytes

    dib_size = 40
    data_offset = 14 + dib_size
    file_size = data_offset + row_stride * image.height

    buf = bytearray()

    # File header (14 bytes)
    buf.extend(b"BM")
    buf.extend(struct.pack("<I", file_size))
    buf.extend(b"\x00\x00\x00\x00")  # reserved
    buf.extend(struct.pack("<I", data_offset))

    # DIB header (40 bytes)
    buf.extend(struct.pack("<I", dib_size))
    buf.extend(struct.pack("<i", image.width))
    buf.extend(struct.pack("<i", image.height))  # positive = bottom-up
    buf.extend(struct.pack("<H", 1))  # planes
    buf.extend(struct.pack("<H", bits_per_pixel))
    buf.extend(struct.pack("<I", 0))  # compression = none
    buf.extend(struct.pack("<I", row_stride * image.height))  # image size
    buf.extend(b"\x00" * 16)  # resolution + palette counts

    # Pixel data (bottom-up)
    pad_bytes = b"\x00" * padding
    for y in range(image.height - 1, -1, -1):
        row_offset = y * image.width * src_bpp
        src_row = image.data[row_offset : row_offset + image.width * src_bpp]
        bmp_row = bytearray(row_data_bytes)
        if bmp_bpp == 3:
            # RGB -> BGR via slice assignment (C-level ops)
            bmp_row[0::3] = src_row[2::3]  # B
            bmp_row[1::3] = src_row[1::3]  # G
            bmp_row[2::3] = src_row[0::3]  # R
        else:
            # RGBA -> BGRA via slice assignment
            bmp_row[0::4] = src_row[2::4]  # B
            bmp_row[1::4] = src_row[1::4]  # G
            bmp_row[2::4] = src_row[0::4]  # R
            bmp_row[3::4] = src_row[3::4]  # A
        buf.extend(bmp_row)
        buf.extend(pad_bytes)

    return _write_dest(dest, bytes(buf))

def convert(image: Image, mode: str) -> Image:
    """Convert an :class:`Image` to a different pixel mode.

    Args:
        image: Source image.
        mode: Target mode (``'L'``, ``'LA'``, ``'RGB'``, ``'RGBA'``).

    Returns:
        New image in the target mode.

    Raises:
        ValueError: If the conversion is not supported.
    """
    if mode not in _MODE_BPP:
        raise ValueError(f"unsupported target mode {mode!r}")
    if image.mode == mode:
        return dataclasses.replace(image)
    if image.bit_depth != 8:
        raise ValueError("convert() only supports 8-bit images")

    src = image.data
    w, h = image.width, image.height
    npx = w * h
    dst_bpp = _MODE_BPP[mode]
    out = bytearray(npx * dst_bpp)

    key = (image.mode, mode)

    if key == ("L", "RGB"):
        out[0::3] = src
        out[1::3] = src
        out[2::3] = src
    elif key == ("L", "RGBA"):
        out[0::4] = src
        out[1::4] = src
        out[2::4] = src
        out[3::4] = b"\xff" * npx
    elif key == ("L", "LA"):
        out[0::2] = src
        out[1::2] = b"\xff" * npx
    elif key == ("LA", "L"):
        out[:] = src[0::2]
    elif key == ("LA", "RGBA"):
        lum = src[0::2]
        out[0::4] = lum
        out[1::4] = lum
        out[2::4] = lum
        out[3::4] = src[1::2]
    elif key == ("LA", "RGB"):
        lum = src[0::2]
        out[0::3] = lum
        out[1::3] = lum
        out[2::3] = lum
    elif key == ("RGB", "RGBA"):
        out[0::4] = src[0::3]
        out[1::4] = src[1::3]
        out[2::4] = src[2::3]
        out[3::4] = b"\xff" * npx
    elif key == ("RGB", "L"):
        for i in range(npx):
            si = i * 3
            out[i] = (
                src[si] * 299 + src[si + 1] * 587 + src[si + 2] * 114 + 500
            ) // 1000
    elif key == ("RGB", "LA"):
        for i in range(npx):
            si = i * 3
            o = i * 2
            out[o] = (
                src[si] * 299 + src[si + 1] * 587 + src[si + 2] * 114 + 500
            ) // 1000
            out[o + 1] = 255
    elif key == ("RGBA", "RGB"):
        out[0::3] = src[0::4]
        out[1::3] = src[1::4]
        out[2::3] = src[2::4]
    elif key == ("RGBA", "L"):
        for i in range(npx):
            si = i * 4
            out[i] = (
                src[si] * 299 + src[si + 1] * 587 + src[si + 2] * 114 + 500
            ) // 1000
    elif key == ("RGBA", "LA"):
        for i in range(npx):
            si = i * 4
            o = i * 2
            out[o] = (
                src[si] * 299 + src[si + 1] * 587 + src[si + 2] * 114 + 500
            ) // 1000
            out[o + 1] = src[si + 3]
    else:
        raise ValueError(f"unsupported conversion {image.mode!r} -> {mode!r}")

    return Image(
        width=w,
        height=h,
        data=bytes(out),
        mode=mode,
        bit_depth=8,
        metadata=image.metadata,
    )

def matrix_to_png(
    matrix: list[list[int | float]] | list[list[Any]],
    dest: str | os.PathLike[str] | IO[bytes] | None = None,
    *,
    bit_depth: int = 8,
    filter_strategy: str = "auto",
    compression_level: int = 6,
) -> bytes | None:
    """Encode a 2-D numeric matrix as a grayscale PNG.

    Integer values in ``[0, 2^bit_depth - 1]`` are stored directly.
    Float values (or out-of-range integers) are linearly mapped to
    the full pixel range; the mapping parameters are stored in a PNG
    ``tEXt`` chunk so :func:`png_to_matrix` can reconstruct them.

    Args:
        matrix: Rows of numeric values (all rows must have equal length).
        dest: File path, file object, or ``None`` to return bytes.
        bit_depth: 8 or 16 bits per sample.
        filter_strategy: PNG row filter strategy.
        compression_level: zlib compression level (0–9).

    Returns:
        PNG bytes when *dest* is ``None``, otherwise ``None``.
    """
    if not matrix or not matrix[0]:
        raise EncodeError("matrix must be non-empty")

    rows = len(matrix)
    cols = len(matrix[0])
    max_val = (1 << bit_depth) - 1

    # Determine dtype and mapping
    all_int = all(isinstance(v, int) for row in matrix for v in row)
    vmin: float = min(min(row) for row in matrix)
    vmax: float = max(max(row) for row in matrix)

    need_mapping = not all_int or vmin < 0 or vmax > max_val
    dtype = "float64" if need_mapping else (f"uint{bit_depth}")

    if need_mapping and vmin == vmax:
        scale = 0.0
    elif need_mapping:
        scale = max_val / (vmax - vmin)
    else:
        scale = 1.0

    # Build pixel data
    if bit_depth == 8:
        pixels = bytearray(rows * cols)
        idx = 0
        if need_mapping:
            for row in matrix:
                for v in row:
                    pixels[idx] = max(0, min(max_val, round((v - vmin) * scale)))
                    idx += 1
        else:
            for row in matrix:
                for v in row:
                    pixels[idx] = int(v)
                    idx += 1
    else:  # 16-bit
        pixels = bytearray(rows * cols * 2)
        idx = 0
        if need_mapping:
            for row in matrix:
                for v in row:
                    pv = max(0, min(max_val, round((v - vmin) * scale)))
                    struct.pack_into(">H", pixels, idx, pv)
                    idx += 2
        else:
            for row in matrix:
                for v in row:
                    struct.pack_into(">H", pixels, idx, int(v))
                    idx += 2

    metadata: dict[str, str] = {}
    meta_info: dict[str, Any] = {
        "dtype": dtype,
        "rows": rows,
        "cols": cols,
    }
    if need_mapping:
        meta_info["min"] = vmin
        meta_info["max"] = vmax
    metadata[_MATRIX_TEXT_KEY] = json.dumps(meta_info, separators=(",", ":"))

    img = Image(
        width=cols,
        height=rows,
        data=bytes(pixels),
        mode="L",
        bit_depth=bit_depth,
        metadata=metadata,
    )
    return encode_png(
        img,
        dest,
        filter_strategy=filter_strategy,
        compression_level=compression_level,
    )

def png_to_matrix(
    source: str | os.PathLike[str] | IO[bytes] | bytes,
) -> list[list[int]] | list[list[float]]:
    """Decode a grayscale PNG into a 2-D numeric matrix.

    If the PNG contains ``zerodep:matrix`` metadata (written by
    :func:`matrix_to_png`), float values are reconstructed from the
    stored scale and offset.  Otherwise raw pixel values are returned
    as integers.

    Args:
        source: File path, file object, or raw PNG bytes.

    Returns:
        2-D list of numeric values.
    """
    img = decode_png(source)

    # Convert to grayscale if needed
    if img.mode != "L":
        img = convert(img, "L")

    max_val = (1 << img.bit_depth) - 1

    # Parse metadata
    meta_info: dict[str, Any] | None = None
    if img.metadata and _MATRIX_TEXT_KEY in img.metadata:
        try:
            meta_info = json.loads(img.metadata[_MATRIX_TEXT_KEY])
        except (json.JSONDecodeError, KeyError):
            meta_info = None

    need_mapping = meta_info is not None and meta_info.get("dtype") == "float64"

    if img.bit_depth == 16:
        # Unpack 16-bit values
        raw_values: list[int] = []
        for i in range(0, len(img.data), 2):
            raw_values.append(struct.unpack(">H", img.data[i : i + 2])[0])
    else:
        raw_values = list(img.data)

    rows = img.height
    cols = img.width

    if need_mapping:
        assert meta_info is not None
        vmin = meta_info["min"]
        vmax = meta_info["max"]
        if vmin == vmax:
            scale_inv = 0.0
        else:
            scale_inv = (vmax - vmin) / max_val
        result_float: list[list[float]] = []
        for y in range(rows):
            row: list[float] = []
            for x in range(cols):
                pv = raw_values[y * cols + x]
                row.append(vmin + pv * scale_inv)
            result_float.append(row)
        return result_float

    result_int: list[list[int]] = []
    for y in range(rows):
        result_int.append(raw_values[y * cols : (y + 1) * cols])
    return result_int