compare_codec/__init__.py

"""Codec registry — protocol, config dataclass, and discovery."""

from __future__ import annotations

from dataclasses import dataclass, field
from pathlib import Path
from typing import Protocol, runtime_checkable

import numpy as np


@dataclass(frozen=True)
class CodecConfig:
    """A single configuration offered by a codec (e.g. '44kHz / 9 quantizers')."""

    name: str
    params: dict = field(default_factory=dict)


@runtime_checkable
class AudioCodec(Protocol):
    """Interface every codec must implement."""

    @property
    def name(self) -> str: ...

    @property
    def sample_rate(self) -> int: ...

    def configs(self) -> list[CodecConfig]: ...

    def encode_decode(self, audio_path: Path, config: CodecConfig) -> np.ndarray:
        """Round-trip: raw file in -> mono float32 numpy array out at self.sample_rate."""
        ...


_REGISTRY: dict[str, AudioCodec] = {}


def register(codec: AudioCodec) -> None:
    _REGISTRY[codec.name] = codec


def get_all() -> dict[str, AudioCodec]:
    return dict(_REGISTRY)


# Import codec modules so they self-register on startup.
from compare_codec import dac as _dac  # noqa: E402, F401
from compare_codec import encodec_codec as _encodec  # noqa: E402, F401
from compare_codec import mimi_codec as _mimi  # noqa: E402, F401
from compare_codec import snac_codec as _snac  # noqa: E402, F401