audio_stream.py

import threading

import numpy as np

try:
    import sounddevice as sd
except ImportError:  # pragma: no cover - optional playback dependency
    sd = None

from typing import Protocol, cast


class _SoundDeviceModule(Protocol):
    OutputStream: type
    CallbackStop: type[BaseException]


class _AudioStreamPlayer:
    def __init__(self, audio_data: np.ndarray, sample_rate: int) -> None:
        if sd is None:
            raise RuntimeError("sounddevice is not installed.")

        if audio_data.ndim != 2:
            raise ValueError("audio_data must be 2D (frames, channels).")

        self._sd = cast(_SoundDeviceModule, cast(object, sd))
        self._audio_data = audio_data.astype(np.float32, copy=False)
        self._sample_rate = int(sample_rate)
        self._position = 0
        self._lock = threading.Lock()
        self._is_playing = False
        self._stream = self._sd.OutputStream(
            samplerate=self._sample_rate,
            channels=self._audio_data.shape[1],
            dtype="float32",
            callback=self._callback,)

    @property
    def is_playing(self) -> bool:
        return self._is_playing

    @property
    def duration_seconds(self) -> float:
        return float(self._audio_data.shape[0]) / float(self._sample_rate)

    def play(self) -> None:
        with self._lock:
            self._is_playing = True
        if not self._stream.active:
            self._stream.start()

    def pause(self) -> None:
        with self._lock:
            self._is_playing = False
        if self._stream.active:
            self._stream.stop()

    def close(self) -> None:
        with self._lock:
            self._is_playing = False
        if self._stream.active:
            self._stream.stop()
        self._stream.close()

    def get_position_seconds(self) -> float:
        with self._lock:
            return float(self._position) / float(self._sample_rate)

    def set_position_seconds(self, seconds: float) -> None:
        clamped = max(0.0, min(seconds, self.duration_seconds))
        with self._lock:
            self._position = int(clamped * self._sample_rate)

    def _callback(self, outdata, frames, _time, _status) -> None:
        with self._lock:
            if not self._is_playing:
                outdata[:] = 0
                return

            start = self._position
            end = start + frames
            chunk = self._audio_data[start:end]
            if chunk.shape[0] < frames:
                outdata[:chunk.shape[0]] = chunk
                outdata[chunk.shape[0]:] = 0
                self._position = 0
                self._is_playing = False
                return

            outdata[:] = chunk
            self._position = end