proc-tap 0.2.0__py3-none-any.whl

proctap/contrib/discord_source.py

@@ -0,0 +1,381 @@
+"""
+Discord AudioSource implementation for proctap.
+
+Streams process audio to Discord voice channels with automatic format conversion
+and resampling to Discord's required format (48kHz, 16-bit PCM, stereo).
+"""
+
+from __future__ import annotations
+
+import logging
+import struct
+import threading
+import time
+from collections import deque
+from typing import Optional
+
+import numpy as np
+
+try:
+    import discord
+except ImportError as e:
+    raise ImportError(
+        "discord.py is required for ProcessAudioSource. "
+        "Install with: pip install discord.py"
+    ) from e
+
+from ..core import ProcessAudioTap
+
+logger = logging.getLogger(__name__)
+
+# Discord audio constants
+DISCORD_SAMPLE_RATE = 48000  # Hz
+DISCORD_CHANNELS = 2  # Stereo
+DISCORD_SAMPLE_SIZE = 2  # 16-bit = 2 bytes
+DISCORD_FRAME_DURATION_MS = 20  # ms
+DISCORD_SAMPLES_PER_FRAME = int(DISCORD_SAMPLE_RATE * DISCORD_FRAME_DURATION_MS / 1000)
+DISCORD_FRAME_SIZE = DISCORD_SAMPLES_PER_FRAME * DISCORD_CHANNELS * DISCORD_SAMPLE_SIZE  # 3840 bytes
+
+
+class ProcessAudioSource(discord.AudioSource):
+    """
+    Discord AudioSource that captures audio from a specific process.
+
+    This class streams audio from a target process to Discord voice channels,
+    automatically handling format detection, conversion, and resampling.
+
+    Args:
+        pid: Process ID to capture audio from
+        gain: Audio gain multiplier (default: 1.0)
+        max_queue_frames: Maximum frames to buffer (default: 50)
+
+    Example:
+        ```python
+        import discord
+        from processaudiotap.contrib import ProcessAudioSource
+
+        # In your Discord bot
+        voice_client = await channel.connect()
+        source = ProcessAudioSource(pid=12345, gain=1.2)
+        voice_client.play(source)
+        ```
+
+    Note:
+        - Automatically resamples from native 44.1kHz to Discord's 48kHz
+        - Handles both 32-bit float and 16-bit PCM input formats
+        - Runs capture in a separate thread for minimal latency
+    """
+
+    def __init__(
+        self,
+        pid: int,
+        gain: float = 1.0,
+        max_queue_frames: int = 50,
+    ) -> None:
+        self.pid = pid
+        self.gain = gain
+        self.max_queue_frames = max_queue_frames
+
+        self._tap: Optional[ProcessAudioTap] = None
+        self._capture_thread: Optional[threading.Thread] = None
+        self._stop_event = threading.Event()
+
+        # Audio queue and buffer
+        self._audio_queue: deque[bytes] = deque(maxlen=max_queue_frames)
+        self._queue_lock = threading.Lock()
+        self._buffer = bytearray()
+
+        # Format detection
+        self._source_format: Optional[dict[str, int]] = None
+        self._is_float32 = False
+        self._format_detected = False
+
+        # Statistics
+        self._frames_dropped = 0
+        self._frames_served = 0
+
+        logger.info(f"ProcessAudioSource created for PID {pid} (gain={gain})")
+
+    def start(self) -> None:
+        """Start audio capture from the target process."""
+        if self._capture_thread is not None:
+            logger.warning("Audio capture already started")
+            return
+
+        logger.info(f"Starting audio capture for PID {self.pid}")
+
+        # Create ProcessAudioTap
+        self._tap = ProcessAudioTap(pid=self.pid)
+        self._tap.start()
+
+        # Get source format
+        self._source_format = self._tap.get_format()
+        logger.info(f"Source format: {self._source_format}")
+
+        # Start capture thread
+        self._stop_event.clear()
+        self._capture_thread = threading.Thread(
+            target=self._capture_loop,
+            daemon=True,
+            name=f"ProcessAudioSource-{self.pid}"
+        )
+        self._capture_thread.start()
+
+        logger.info("Audio capture started")
+
+    def stop(self) -> None:
+        """Stop audio capture and release resources."""
+        if self._capture_thread is None:
+            return
+
+        logger.info("Stopping audio capture...")
+        self._stop_event.set()
+
+        if self._capture_thread.is_alive():
+            self._capture_thread.join(timeout=2.0)
+            self._capture_thread = None
+
+        if self._tap is not None:
+            try:
+                self._tap.close()
+            except Exception:
+                logger.exception("Error closing ProcessAudioTap")
+            finally:
+                self._tap = None
+
+        logger.info(
+            f"Audio capture stopped. Stats: served={self._frames_served}, dropped={self._frames_dropped}"
+        )
+
+    def _capture_loop(self) -> None:
+        """
+        Capture loop running in separate thread.
+        Continuously reads audio from ProcessAudioTap and queues converted frames.
+        """
+        logger.debug("Capture loop started")
+
+        while not self._stop_event.is_set():
+            try:
+                # Read audio with timeout
+                chunk = self._tap.read(timeout=0.5)
+
+                if chunk is None or len(chunk) == 0:
+                    continue
+
+                # Detect format on first chunk
+                if not self._format_detected:
+                    self._detect_format(chunk)
+
+                # Convert and resample audio
+                converted = self._convert_audio(chunk)
+
+                if converted:
+                    with self._queue_lock:
+                        try:
+                            self._audio_queue.append(converted)
+                        except IndexError:
+                            # Queue full, frame dropped
+                            self._frames_dropped += 1
+
+            except Exception:
+                logger.exception("Error in capture loop")
+                time.sleep(0.1)  # Prevent busy loop on error
+
+        logger.debug("Capture loop ended")
+
+    def _detect_format(self, chunk: bytes) -> None:
+        """
+        Detect if audio data is 32-bit float or 16-bit PCM.
+
+        WASAPI may return 32-bit float despite requesting 16-bit PCM.
+        """
+        if self._format_detected:
+            return
+
+        if len(chunk) < 8:
+            return  # Not enough data to detect
+
+        # Try to interpret as 32-bit float
+        try:
+            sample_count = len(chunk) // 4
+            floats = struct.unpack(f"{sample_count}f", chunk[:sample_count * 4])
+
+            # Check if values are in typical float range [-1.0, 1.0]
+            max_abs = max(abs(f) for f in floats[:min(100, len(floats))])
+
+            if 0.0 < max_abs <= 2.0:  # Likely float32
+                self._is_float32 = True
+                logger.info("Detected 32-bit float PCM format")
+            else:
+                self._is_float32 = False
+                logger.info("Detected 16-bit PCM format")
+
+        except struct.error:
+            self._is_float32 = False
+            logger.info("Defaulting to 16-bit PCM format")
+
+        self._format_detected = True
+
+    def _convert_audio(self, chunk: bytes) -> Optional[bytes]:
+        """
+        Convert audio to Discord format (48kHz, 16-bit PCM, stereo).
+
+        Args:
+            chunk: Raw audio data from ProcessAudioTap
+
+        Returns:
+            Converted audio data, or None if conversion failed
+        """
+        if not self._source_format:
+            return None
+
+        try:
+            # Parse audio data
+            if self._is_float32:
+                # 32-bit float PCM
+                sample_count = len(chunk) // 4
+                audio_data = np.frombuffer(chunk, dtype=np.float32, count=sample_count)
+
+                # Check for NaN/Inf
+                if np.any(~np.isfinite(audio_data)):
+                    logger.warning("NaN/Inf detected in audio data, skipping frame")
+                    return None
+
+                # Convert float32 [-1.0, 1.0] to int16
+                audio_data = np.clip(audio_data * 32767.0 * self.gain, -32768, 32767)
+                audio_data = audio_data.astype(np.int16)
+            else:
+                # 16-bit PCM
+                sample_count = len(chunk) // 2
+                audio_data = np.frombuffer(chunk, dtype=np.int16, count=sample_count)
+
+                # Apply gain
+                if self.gain != 1.0:
+                    audio_data = audio_data.astype(np.float32)
+                    audio_data = np.clip(audio_data * self.gain, -32768, 32767)
+                    audio_data = audio_data.astype(np.int16)
+
+            # Reshape to (frames, channels)
+            channels = self._source_format["channels"]
+            frames = len(audio_data) // channels
+
+            if frames == 0:
+                return None
+
+            audio_data = audio_data[:frames * channels].reshape(frames, channels)
+
+            # Handle mono to stereo conversion
+            if channels == 1:
+                audio_data = np.repeat(audio_data, 2, axis=1)
+
+            # Resample if necessary
+            source_rate = self._source_format["sample_rate"]
+            if source_rate != DISCORD_SAMPLE_RATE:
+                audio_data = self._resample(audio_data, source_rate, DISCORD_SAMPLE_RATE)
+
+            # Convert to bytes
+            return audio_data.tobytes()
+
+        except Exception:
+            logger.exception("Error converting audio")
+            return None
+
+    def _resample(
+        self,
+        audio: np.ndarray,
+        source_rate: int,
+        target_rate: int
+    ) -> np.ndarray:
+        """
+        Resample audio using linear interpolation.
+
+        Args:
+            audio: Audio data shaped (frames, channels)
+            source_rate: Source sample rate
+            target_rate: Target sample rate
+
+        Returns:
+            Resampled audio data
+        """
+        if source_rate == target_rate:
+            return audio
+
+        # Calculate resampling ratio
+        ratio = target_rate / source_rate
+        source_frames = len(audio)
+        target_frames = int(source_frames * ratio)
+
+        # Linear interpolation for each channel
+        source_indices = np.arange(source_frames)
+        target_indices = np.linspace(0, source_frames - 1, target_frames)
+
+        resampled = np.empty((target_frames, audio.shape[1]), dtype=audio.dtype)
+
+        for ch in range(audio.shape[1]):
+            resampled[:, ch] = np.interp(target_indices, source_indices, audio[:, ch])
+
+        return resampled
+
+    def read(self) -> bytes:
+        """
+        Read one Discord audio frame (20ms @ 48kHz = 3840 bytes).
+
+        This method is called by discord.py's voice client.
+
+        Returns:
+            3840 bytes of 16-bit PCM stereo audio, or silence if no data available
+        """
+        # Accumulate data until we have a full Discord frame
+        while len(self._buffer) < DISCORD_FRAME_SIZE:
+            with self._queue_lock:
+                if not self._audio_queue:
+                    # No data available, return silence
+                    silence = b"\x00" * DISCORD_FRAME_SIZE
+                    return silence
+
+                chunk = self._audio_queue.popleft()
+                self._buffer.extend(chunk)
+
+        # Extract one Discord frame
+        frame = bytes(self._buffer[:DISCORD_FRAME_SIZE])
+        del self._buffer[:DISCORD_FRAME_SIZE]
+
+        self._frames_served += 1
+        return frame
+
+    def is_opus(self) -> bool:
+        """
+        Indicate whether this source provides Opus-encoded audio.
+
+        Returns:
+            False (this source provides raw PCM)
+        """
+        return False
+
+    def cleanup(self) -> None:
+        """
+        Cleanup resources when discord.py is done with this source.
+
+        This is called automatically by discord.py.
+        """
+        self.stop()
+
+    @property
+    def stats(self) -> dict[str, int]:
+        """
+        Get capture statistics.
+
+        Returns:
+            Dictionary with keys:
+            - 'frames_served': Number of frames successfully served
+            - 'frames_dropped': Number of frames dropped due to queue overflow
+            - 'queue_size': Current number of frames in queue
+        """
+        with self._queue_lock:
+            queue_size = len(self._audio_queue)
+
+        return {
+            "frames_served": self._frames_served,
+            "frames_dropped": self._frames_dropped,
+            "queue_size": queue_size,
+        }

proctap/core.py

@@ -0,0 +1,222 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Callable, Optional, AsyncIterator
+import threading
+import queue
+import asyncio
+import logging
+
+logger = logging.getLogger(__name__)
+
+# -------------------------------
+# Backend import (platform-specific)
+# -------------------------------
+
+from .backends import get_backend
+from .backends.base import AudioBackend
+
+AudioCallback = Callable[[bytes, int], None]  # (pcm_bytes, num_frames)
+
+
+@dataclass
+class StreamConfig:
+    sample_rate: int = 44100   # Hz
+    channels: int = 2
+    # NOTE:
+    # 現状 backend 側でバッファサイズは制御していないので
+    # frames_per_buffer は「論理的なサイズ」として扱うだけ。
+    frames_per_buffer: int = 480  # 10ms @ 48kHz
+
+
+class ProcessAudioTap:
+    """
+    High-level API for process-specific audio capture.
+
+    Supports multiple platforms:
+    - Windows: WASAPI Process Loopback (fully implemented)
+    - Linux: PulseAudio/PipeWire (under development)
+    - macOS: Core Audio (planned, not yet implemented)
+
+    Usage:
+    - Callback mode: start(on_data=callback)
+    - Async mode: async for chunk in tap.iter_chunks()
+    """
+
+    def __init__(
+        self,
+        pid: int,
+        config: StreamConfig | None = None,
+        on_data: Optional[AudioCallback] = None,
+    ) -> None:
+        self._pid = pid
+        self._cfg = config or StreamConfig()
+        self._on_data = on_data
+
+        # Get platform-specific backend
+        # Pass config to backend (Windows uses fixed format, Linux uses config)
+        self._backend: AudioBackend = get_backend(
+            pid=pid,
+            sample_rate=self._cfg.sample_rate,
+            channels=self._cfg.channels,
+            sample_width=2,  # 16-bit = 2 bytes
+        )
+        logger.debug(f"Using backend: {type(self._backend).__name__}")
+
+        self._thread: Optional[threading.Thread] = None
+        self._stop_event = threading.Event()
+        self._async_queue: "queue.Queue[bytes | None]" = queue.Queue()
+
+    # --- public API -----------------------------------------------------
+
+    def start(self) -> None:
+        if self._thread is not None:
+            # すでに start 済みなら何もしない
+            return
+
+        # Start platform-specific backend
+        self._backend.start()
+
+        self._stop_event.clear()
+        self._thread = threading.Thread(target=self._worker, daemon=True)
+        self._thread.start()
+
+    def stop(self) -> None:
+        self._stop_event.set()
+
+        if self._thread is not None:
+            self._thread.join(timeout=1.0)
+            self._thread = None
+
+        try:
+            self._backend.stop()
+        except Exception:
+            logger.exception("Error while stopping capture")
+
+    def close(self) -> None:
+        self.stop()
+
+    def __enter__(self) -> "ProcessAudioTap":
+        self.start()
+        return self
+
+    def __exit__(self, exc_type, exc, tb) -> None:
+        self.close()
+
+    # --- properties -----------------------------------------------------
+
+    @property
+    def is_running(self) -> bool:
+        """Check if audio capture is currently running."""
+        return self._thread is not None and self._thread.is_alive()
+
+    @property
+    def pid(self) -> int:
+        """Get the target process ID."""
+        return self._pid
+
+    @property
+    def config(self) -> StreamConfig:
+        """Get the stream configuration (note: does not affect native backend)."""
+        return self._cfg
+
+    # --- utility methods ------------------------------------------------
+
+    def set_callback(self, callback: Optional[AudioCallback]) -> None:
+        """
+        Change the audio data callback.
+
+        Args:
+            callback: New callback function, or None to remove callback
+        """
+        self._on_data = callback
+
+    def get_format(self) -> dict[str, int]:
+        """
+        Get audio format information from the native backend.
+
+        Returns:
+            Dictionary with keys:
+            - 'sample_rate': Sample rate in Hz (e.g., 44100)
+            - 'channels': Number of channels (e.g., 2 for stereo)
+            - 'bits_per_sample': Bits per sample (e.g., 16)
+        """
+        return self._backend.get_format()
+
+    def read(self, timeout: float = 1.0) -> Optional[bytes]:
+        """
+        Synchronous API: Read one audio chunk (blocking).
+
+        Args:
+            timeout: Maximum time to wait for data in seconds
+
+        Returns:
+            PCM audio data as bytes, or None if timeout or no data
+
+        Note:
+            This is a simple synchronous alternative to the async API.
+            The capture must be started first with start().
+        """
+        if not self.is_running:
+            raise RuntimeError("Capture is not running. Call start() first.")
+
+        try:
+            return self._async_queue.get(timeout=timeout)
+        except queue.Empty:
+            return None
+
+    # --- async interface ------------------------------------------------
+
+    async def iter_chunks(self) -> AsyncIterator[bytes]:
+        """
+        Async generator that yields PCM chunks as bytes.
+        """
+        loop = asyncio.get_running_loop()
+
+        while True:
+            chunk = await loop.run_in_executor(None, self._async_queue.get)
+            if chunk is None:  # sentinel
+                break
+            yield chunk
+
+    # --- worker thread --------------------------------------------------
+
+    def _worker(self) -> None:
+        """
+        Loop:
+            data = backend.read()
+            -> callback
+            -> async_queue
+        """
+        while not self._stop_event.is_set():
+            try:
+                data = self._backend.read()
+            except Exception:
+                logger.exception("Error reading data from backend")
+                continue
+
+            if not data:
+                # パケットがまだ無いケース。ここで sleep 入れるかは後で調整。
+                continue
+
+            # callback
+            if self._on_data is not None:
+                try:
+                    # frames 数は backend から直接取れないので、とりあえず -1 を渡す。
+                    # TODO: _backend.get_format() を見て frame 数を計算する改善余地あり。
+                    self._on_data(data, -1)
+                except Exception:
+                    logger.exception("Error in audio callback")
+
+            # async queue
+            try:
+                self._async_queue.put_nowait(data)
+            except queue.Full:
+                # リアルタイム性重視なので捨てる
+                pass
+
+        # 終了シグナル
+        try:
+            self._async_queue.put_nowait(None)
+        except queue.Full:
+            pass