audio_capture.py

"""Audio device enumeration and capture through SoundCard/WASAPI."""

from __future__ import annotations

from dataclasses import dataclass
from queue import Empty, Full, Queue
from threading import Event, Thread
from typing import Any

import numpy as np


def _soundcard_backend():
    """Load SoundCard after QApplication has initialized Windows COM."""
    import soundcard as sc

    return sc


DISABLED_DEVICE = "__none__"


@dataclass(frozen=True)
class AudioDevice:
    id: str
    name: str
    source: str


class AudioCapture:
    """Capture output loopback and input independently."""

    def __init__(
        self,
        output: dict[str, Any],
        input_config: dict[str, Any],
    ) -> None:
        self.output = output
        self.input = input_config
        self._blocks: dict[str, Queue[np.ndarray]] = {
            "output": Queue(maxsize=2),
            "input": Queue(maxsize=2),
        }
        self._stop = Event()
        self._threads: list[Thread] = []
        self.errors: dict[str, Exception] = {}

    @staticmethod
    def list_devices() -> list[AudioDevice]:
        sc = _soundcard_backend()
        devices = [
            AudioDevice(device.id, device.name, "input")
            for device in sc.all_microphones(include_loopback=False)
        ]
        devices.extend(
            AudioDevice(device.id, device.name, "system")
            for device in sc.all_microphones(include_loopback=True)
            if getattr(device, "isloopback", False)
        )
        return devices

    def start(self) -> None:
        if any(thread.is_alive() for thread in self._threads):
            return
        self._stop.clear()
        self.errors.clear()
        self._threads = []
        for stream in ("output", "input"):
            thread = Thread(target=self._capture_loop, args=(stream,), name=f"audio-{stream}", daemon=True)
            self._threads.append(thread)
            thread.start()

    def stop(self) -> None:
        self._stop.set()
        for thread in self._threads:
            thread.join(timeout=2)
        self._threads = []

    def read_latest(self) -> dict[str, np.ndarray | None]:
        latest: dict[str, np.ndarray | None] = {}
        for stream, blocks in self._blocks.items():
            latest[stream] = None
            try:
                while True:
                    latest[stream] = blocks.get_nowait()
            except Empty:
                pass
        return latest

    def _select_device(self, stream: str) -> Any:
        sc = _soundcard_backend()
        settings = self.output if stream == "output" else self.input
        device_id = settings["source"]
        if device_id == DISABLED_DEVICE:
            return None
        if device_id:
            for device in self.list_devices():
                source = "system" if stream == "output" else "input"
                if device.id == device_id and device.source == source:
                    return sc.get_microphone(device.id, include_loopback=stream == "output")
        if stream == "output":
            speaker = sc.default_speaker()
            return sc.get_microphone(speaker.id, include_loopback=True)
        return sc.default_microphone()

    def _capture_loop(self, stream: str) -> None:
        try:
            settings = self.output if stream == "output" else self.input
            device = self._select_device(stream)
            if device is None:
                return
            channels = 1 if stream == "output" and settings["channels"] == "mono" else 2
            with device.recorder(
                samplerate=settings["sample_rate"],
                channels=channels,
                blocksize=settings["block_size"],
            ) as recorder:
                while not self._stop.is_set():
                    block = np.asarray(recorder.record(numframes=settings["block_size"]), dtype=np.float32)
                    if block.ndim == 2 and channels == 1:
                        block = np.mean(block, axis=1)
                    try:
                        self._blocks[stream].put_nowait(block)
                    except Full:
                        try:
                            self._blocks[stream].get_nowait()
                        except Empty:
                            pass
                        self._blocks[stream].put_nowait(block)
        except Exception as error:
            self.errors[stream] = error


def print_devices() -> None:
    """Print devices for a quick Windows audio-backend diagnostic."""
    for device in AudioCapture.list_devices():
        print(f"{device.source}: {device.id} - {device.name}")


if __name__ == "__main__":
    print_devices()