package/container.py - master - gekkoff/tkisi_task_2

package/container.py
144 строки4 KB
Alex
Пришлось пересоздать репозиторий, поэтому мало коммитов
09 янв 2026, 23:53
09 янв 2026, 23:53ff15d6e
"""
Формат (Little Endian):
- magic: 4 bytes = "ICF1"
- version: uint8
- algorithm_id: uint8 (1=RLE, 2=LZW, 3=FRACTAL)
- reserved: uint16 (0)
- width: uint32
- height: uint32
- bpp: uint16 (в проекте всегда 24)
- meta_len: uint16
- meta: meta_len bytes
- payload: остальные байты
"""

from __future__ import annotations

import struct
from dataclasses import dataclass

from .algorithm import Algorithm


MAGIC = b"ICF1"
VERSION = 1
BPP_24 = 24


class ContainerError(ValueError):
    """Ошибка контейнера (повреждённый или неподдерживаемый формат)."""


_ALG_ID = {
    Algorithm.RLE: 1,
    Algorithm.LZW: 2,
    Algorithm.FRACTAL: 3,
}
_ID_ALG = {v: k for k, v in _ALG_ID.items()}


@dataclass(frozen=True)
class Container:
    """
    Распарсенный контейнер.

    Attributes:
        algorithm: алгоритм
        width: ширина
        height: высота
        bpp: бит на пиксель (24)
        meta: метаданные алгоритма (может быть пусто)
        payload: сжатые данные
    """
    algorithm: Algorithm
    width: int
    height: int
    bpp: int
    meta: bytes
    payload: bytes


def encode_container(algorithm: Algorithm, width: int, height: int, bpp: int, meta: bytes, payload: bytes) -> bytes:
    """
    Сформировать байты контейнера.

    Args:
        algorithm: выбранный алгоритм
        width: ширина
        height: высота
        bpp: биты на пиксель (должно быть 24)
        meta: метаданные алгоритма
        payload: полезная нагрузка (сжатые байты)

    Returns:
        bytes: готовый контейнер.

    Raises:
        ValueError: если параметры некорректны.
    """
    if bpp != BPP_24:
        raise ValueError("Поддерживается только 24 bpp.")
    if width <= 0 or height <= 0:
        raise ValueError("Некорректные размеры изображения.")
    if len(meta) > 65535:
        raise ValueError("meta слишком большая (max 65535).")

    alg_id = _ALG_ID[algorithm]
    header = struct.pack(
        "<4sBBHIIHH",
        MAGIC,
        VERSION,
        alg_id,
        0,          # reserved
        width,
        height,
        bpp,
        len(meta),
    )
    return header + meta + payload


def decode_container(blob: bytes) -> Container:
    """
    Распарсить контейнер из bytes.

    Args:
        blob: содержимое файла целиком.

    Returns:
        Container: структура с данными.

    Raises:
        ContainerError: если формат неверный.
    """
    header_size = struct.calcsize("<4sBBHIIHH")
    if len(blob) < header_size:
        raise ContainerError("Контейнер слишком маленький.")

    magic, version, alg_id, reserved, width, height, bpp, meta_len = struct.unpack_from("<4sBBHIIHH", blob, 0)

    if magic != MAGIC:
        raise ContainerError("Неверная сигнатура контейнера.")
    if version != VERSION:
        raise ContainerError("Неподдерживаемая версия контейнера.")
    if alg_id not in _ID_ALG:
        raise ContainerError("Неизвестный algorithm_id.")
    if bpp != BPP_24:
        raise ContainerError("Неподдерживаемый bpp (нужен 24).")

    meta_start = header_size
    meta_end = meta_start + meta_len
    if meta_end > len(blob):
        raise ContainerError("Некорректное meta_len (выход за границы).")

    meta = blob[meta_start:meta_end]
    payload = blob[meta_end:]

    return Container(
        algorithm=_ID_ALG[alg_id],
        width=width,
        height=height,
        bpp=bpp,
        meta=meta,
        payload=payload,
    )