converter.py

from __future__ import annotations

import io
import logging
import time
from dataclasses import dataclass
from typing import Optional

import pillow_avif  # noqa: F401  registers AVIF support with Pillow
from PIL import Image, UnidentifiedImageError

from config import Config


logger = logging.getLogger("image_processor")


class ConversionError(Exception):
    """Raised when an image cannot be decoded or converted."""


@dataclass
class ConversionResult:
    image: Image.Image
    source_format: str
    elapsed_seconds: float


class Converter:
    """Decode arbitrary input bytes and normalize to an RGB Pillow image."""

    def __init__(self, config: Config) -> None:
        self.config = config

    def convert(self, raw_bytes: bytes) -> ConversionResult:
        start = time.perf_counter()
        try:
            with Image.open(io.BytesIO(raw_bytes)) as img:
                img.load()
                source_format = (img.format or "UNKNOWN").lower()
                rgb_image = self._flatten_to_rgb(img)
        except UnidentifiedImageError as exc:
            raise ConversionError(f"Could not identify image: {exc}") from exc
        except OSError as exc:
            raise ConversionError(f"Corrupted image data: {exc}") from exc

        elapsed = time.perf_counter() - start
        return ConversionResult(
            image=rgb_image,
            source_format=source_format,
            elapsed_seconds=elapsed,
        )

    def _flatten_to_rgb(self, img: Image.Image) -> Image.Image:
        """Convert any image to RGB, flattening transparency over a solid background."""
        if img.mode == "RGB":
            return img.copy()

        if img.mode in ("RGBA", "LA") or (img.mode == "P" and "transparency" in img.info):
            rgba = img.convert("RGBA")
            background = Image.new("RGB", rgba.size, self.config.background_color)
            background.paste(rgba, mask=rgba.split()[-1])
            return background

        return img.convert("RGB")