zynko 0.1.0__py3-none-any.whl

zynko/player.py

@@ -0,0 +1,413 @@
+"""Plays audio into a voice connection.
+
+Files and links through ffmpeg (any format and stream it opens: mp3, ogg, wav, radio
+streams...), encoded streams, raw PCM, or tones. A queue, pause, skip, volume. Audio goes out
+as 48 kHz stereo 16-bit, in 20 ms frames, paced in real time by the voice connection.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import inspect
+import math
+import os
+import re
+import sys
+from array import array
+from collections.abc import AsyncGenerator, AsyncIterable, AsyncIterator, Callable, Iterable
+from dataclasses import dataclass, field
+from typing import Any, Generic, Protocol, TypeVar, Union, overload
+
+from ._events import Emitter
+
+__all__ = [
+    "SAMPLE_RATE",
+    "CHANNELS",
+    "FRAME_SAMPLES",
+    "FRAME_BYTES",
+    "PCM",
+    "Stream",
+    "AudioInput",
+    "Track",
+    "AudioPlayer",
+    "FrameSink",
+    "tone",
+    "scale_frame",
+    "frames",
+]
+
+SAMPLE_RATE = 48_000
+CHANNELS = 2
+#: 20 ms per channel.
+FRAME_SAMPLES = 960
+FRAME_BYTES = FRAME_SAMPLES * CHANNELS * 2
+
+_CHUNK = 64 * 1024
+
+#: What a byte source may be: bytes, an (async) iterable of bytes, or a file-like object
+#: with ``read`` (a file, ``sys.stdin.buffer``, an ``asyncio.StreamReader``).
+ByteSource = Union[bytes, bytearray, memoryview, Iterable[bytes], AsyncIterable[bytes], Any]
+
+
+@dataclass(slots=True)
+class PCM:
+    """Raw PCM already at 48 kHz, stereo, 16-bit little-endian. Not decoded by ffmpeg."""
+
+    source: ByteSource
+
+
+@dataclass(slots=True)
+class Stream:
+    """Encoded audio (an mp3 download, a pipe), decoded by ffmpeg."""
+
+    source: ByteSource
+
+
+#: What to play: a path or an http(s) link (opened by ffmpeg), a :class:`Stream` or a :class:`PCM`.
+AudioInput = Union[str, "os.PathLike[str]", PCM, Stream]
+
+M = TypeVar("M")
+F = TypeVar("F", bound=Callable[..., Any])
+
+
+@dataclass(slots=True)
+class Track(Generic[M]):
+    input: AudioInput
+    #: Whatever you want back in events: a title, who asked for it...
+    meta: M | None = None
+
+
+class FrameSink(Protocol):
+    """Where the player writes frames: the voice connection."""
+
+    async def capture(self, frame: bytes) -> None: ...
+
+    def flush(self) -> None:
+        """Drops what is buffered, so a skip or pause is heard at once."""
+
+
+def tone(hz: float = 440, seconds: float = 1, volume: float = 0.2) -> PCM:
+    """A sine wave, for testing: ``player.play(tone(440, 2))``."""
+
+    def gen() -> Iterable[bytes]:
+        total = round(seconds * SAMPLE_RATE)
+        step = 2 * math.pi * hz / SAMPLE_RATE
+        amp = volume * 32767
+        for start in range(0, total, FRAME_SAMPLES):
+            n = min(FRAME_SAMPLES, total - start)
+            buf = array("h", bytes(n * CHANNELS * 2))
+            for i in range(n):
+                v = round(math.sin(step * (start + i)) * amp)
+                buf[i * 2] = v
+                buf[i * 2 + 1] = v
+            if sys.byteorder == "big":
+                buf.byteswap()
+            yield buf.tobytes()
+
+    return PCM(gen())
+
+
+def scale_frame(frame: bytes, volume: float) -> bytes:
+    """Multiplies 16-bit little-endian samples by ``volume``, clipping."""
+    a = array("h")
+    a.frombytes(frame)
+    if sys.byteorder == "big":
+        a.byteswap()
+    out = array("h", [max(-32768, min(32767, round(s * volume))) for s in a])
+    if sys.byteorder == "big":
+        out.byteswap()
+    return out.tobytes()
+
+
+def frames(carry: bytes, chunk: bytes) -> tuple[list[bytes], bytes]:
+    """Cuts ``carry + chunk`` into whole 20 ms frames; returns them and what is left over."""
+    data = carry + chunk if carry else bytes(chunk)
+    n = len(data) // FRAME_BYTES * FRAME_BYTES
+    return [data[i : i + FRAME_BYTES] for i in range(0, n, FRAME_BYTES)], data[n:]
+
+
+async def _iter_bytes(src: ByteSource) -> AsyncIterator[bytes]:
+    """Any byte source, as an async iterator of chunks."""
+    if isinstance(src, (bytes, bytearray, memoryview)):
+        mv = memoryview(src).cast("B")
+        for i in range(0, len(mv), _CHUNK):
+            yield bytes(mv[i : i + _CHUNK])
+        return
+    if hasattr(src, "__aiter__"):
+        async for chunk in src:
+            yield bytes(chunk)
+        return
+    read = getattr(src, "read", None)
+    if read is not None:
+        if inspect.iscoroutinefunction(read):
+            while chunk := await read(_CHUNK):
+                yield bytes(chunk)
+        else:
+            while chunk := await asyncio.to_thread(read, _CHUNK):
+                yield bytes(chunk)
+        return
+    for chunk in src:
+        yield bytes(chunk)
+        await asyncio.sleep(0)
+
+
+class _Run:
+    """One track being played: its processes, and whether it was cut short."""
+
+    def __init__(self) -> None:
+        self.stopped = False
+        self.procs: list[asyncio.subprocess.Process] = []
+        self.tasks: list[asyncio.Task[Any]] = []
+
+    def abort(self) -> None:
+        self.stopped = True
+        self.kill()
+
+    def kill(self) -> None:
+        for p in self.procs:
+            if p.returncode is None:
+                try:
+                    p.kill()
+                except (ProcessLookupError, OSError):
+                    pass
+        for t in self.tasks:
+            t.cancel()
+
+    async def reap(self) -> None:
+        """Waits for the killed processes, so their pipes close cleanly."""
+        for p in self.procs:
+            try:
+                # Reading what is left lets the pipe reach its end and close.
+                if p.stdout is not None:
+                    await asyncio.wait_for(_drain(p.stdout), 5)
+                await asyncio.wait_for(p.wait(), 5)
+            except (asyncio.TimeoutError, ProcessLookupError, OSError):
+                pass
+        if self.tasks:
+            await asyncio.gather(*self.tasks, return_exceptions=True)
+
+
+class AudioPlayer(Generic[M]):
+    """A queue of tracks played one after another.
+
+    Events: ``start(track)``, ``finish(track)``, ``error(error, track)`` and ``idle()`` (nothing
+    left in the queue).
+    """
+
+    def __init__(self, sink: FrameSink, *, ffmpeg: str | None = None, ytdlp: bool | str = False) -> None:
+        self._sink = sink
+        #: ffmpeg's path. Default: FFMPEG_PATH, or ``ffmpeg`` on the PATH.
+        self.ffmpeg = ffmpeg or os.environ.get("FFMPEG_PATH") or "ffmpeg"
+        #: Hand http(s) links to yt-dlp first (its path, or True for ``yt-dlp`` on the PATH).
+        self.ytdlp = ytdlp
+        #: Played after the current one, in order.
+        self.queue: list[Track[M]] = []
+        #: What is playing now.
+        self.current: Track[M] | None = None
+        self.paused = False
+        self._volume = 1.0
+        self._events = Emitter("zynko.player")
+        self._run: _Run | None = None
+        self._resumed = asyncio.Event()
+        self._resumed.set()
+        self._task: asyncio.Task[None] | None = None
+
+    # ---------- events ----------
+
+    @overload
+    def on(self, event: str) -> Callable[[F], F]: ...
+    @overload
+    def on(self, event: str, handler: F) -> F: ...
+    def on(self, event: str, handler: F | None = None) -> Any:
+        return self._events.on(event, handler) if handler is not None else self._events.on(event)
+
+    def off(self, event: str, handler: Callable[..., Any]) -> None:
+        self._events.off(event, handler)
+
+    async def wait_for(self, event: str, check: Callable[..., bool] | None = None, timeout: float | None = None) -> Any:
+        return await self._events.wait_for(event, check, timeout)
+
+    # ---------- controls ----------
+
+    @property
+    def volume(self) -> float:
+        """0 (silent) to 2 (twice as loud); 1 is as it came."""
+        return self._volume
+
+    @volume.setter
+    def volume(self, v: float) -> None:
+        self._volume = max(0.0, min(2.0, float(v)))
+
+    @property
+    def playing(self) -> bool:
+        return self._task is not None and not self._task.done()
+
+    def play(self, input: AudioInput, meta: M | None = None) -> Track[M]:
+        """Plays now, in place of what was playing (the queue stays)."""
+        track: Track[M] = Track(input, meta)
+        self.queue.insert(0, track)
+        if self.playing:
+            self.skip()
+        else:
+            self._start()
+        return track
+
+    def enqueue(self, input: AudioInput, meta: M | None = None) -> Track[M]:
+        """Adds to the end of the queue; starts if nothing is playing."""
+        track: Track[M] = Track(input, meta)
+        self.queue.append(track)
+        if not self.playing:
+            self._start()
+        return track
+
+    def skip(self) -> None:
+        """Goes to the next in the queue."""
+        if self._run is not None:
+            self._run.abort()
+        self._sink.flush()
+        self.resume()
+
+    def stop(self) -> None:
+        """Stops and empties the queue."""
+        self.queue.clear()
+        self.skip()
+
+    def pause(self) -> None:
+        if self.paused or self.current is None:
+            return
+        self.paused = True
+        self._resumed.clear()
+        self._sink.flush()
+
+    def resume(self) -> None:
+        self.paused = False
+        self._resumed.set()
+
+    # ---------- the loop ----------
+
+    def _start(self) -> None:
+        self._task = asyncio.ensure_future(self._loop())
+
+    async def _loop(self) -> None:
+        try:
+            while self.queue:
+                track = self.queue.pop(0)
+                self.current = track
+                self._events.emit("start", track)
+                try:
+                    await self._play_one(track)
+                    self._events.emit("finish", track)
+                except asyncio.CancelledError:
+                    raise
+                except Exception as e:  # noqa: BLE001 - reported as an event, the queue goes on
+                    self._events.emit("error", e, track)
+        finally:
+            self.current = None
+            self.paused = False
+            self._resumed.set()
+            self._events.emit("idle")
+
+    async def _play_one(self, track: Track[M]) -> None:
+        run = _Run()
+        self._run = run
+        carry = b""
+        pcm = self._open(track.input, run)
+        try:
+            async for chunk in pcm:
+                if run.stopped:
+                    break
+                out, carry = frames(carry, chunk)
+                for frame in out:
+                    while self.paused and not run.stopped:
+                        await self._resumed.wait()
+                    if run.stopped:
+                        break
+                    await self._send(frame)
+                if run.stopped:
+                    break
+            # The last partial frame, padded with silence.
+            if not run.stopped and carry:
+                await self._send(carry[: len(carry) - len(carry) % 2].ljust(FRAME_BYTES, b"\0"))
+            if not run.stopped:
+                for p in run.procs:
+                    code = await p.wait()
+                    if code != 0 and not run.stopped:
+                        raise RuntimeError(f"{getattr(p, '_zynko_name', 'processo')} terminou com código {code}")
+        finally:
+            if self._run is run:
+                self._run = None
+            run.kill()
+            await pcm.aclose()
+            await run.reap()
+
+    async def _send(self, frame: bytes) -> None:
+        if self._volume != 1:
+            frame = scale_frame(frame, self._volume)
+        await self._sink.capture(frame)
+
+    async def _spawn(self, run: _Run, name: str, args: list[str], stdin: bool) -> asyncio.subprocess.Process:
+        try:
+            p = await asyncio.create_subprocess_exec(
+                *args,
+                stdin=asyncio.subprocess.PIPE if stdin else asyncio.subprocess.DEVNULL,
+                stdout=asyncio.subprocess.PIPE,
+                stderr=None,
+            )
+        except FileNotFoundError:
+            raise RuntimeError(f"{name} não encontrado: instale-o ou indique o caminho") from None
+        p._zynko_name = name  # type: ignore[attr-defined]
+        run.procs.append(p)
+        return p
+
+    async def _open(self, input: AudioInput, run: _Run) -> AsyncGenerator[bytes, None]:
+        """The input as raw PCM chunks."""
+        if isinstance(input, PCM):
+            async for chunk in _iter_bytes(input.source):
+                yield chunk
+            return
+        out = ["-f", "s16le", "-ar", str(SAMPLE_RATE), "-ac", str(CHANNELS), "-loglevel", "error", "pipe:1"]
+        feed: ByteSource | asyncio.StreamReader | None = None
+        if isinstance(input, Stream):
+            feed = input.source
+            args = ["-i", "pipe:0", *out]
+        else:
+            target = os.fspath(input)
+            http = re.match(r"^https?://", target, re.I) is not None
+            if http and self.ytdlp:
+                ytdlp = self.ytdlp if isinstance(self.ytdlp, str) else "yt-dlp"
+                y = await self._spawn(
+                    run, "yt-dlp", [ytdlp, "-f", "bestaudio/best", "--no-playlist", "--quiet", "-o", "-", target], False
+                )
+                assert y.stdout is not None
+                feed = y.stdout
+                args = ["-i", "pipe:0", *out]
+            else:
+                reconnect = ["-reconnect", "1", "-reconnect_streamed", "1", "-reconnect_delay_max", "5"] if http else []
+                args = [*reconnect, "-i", target, *out]
+        f = await self._spawn(run, "ffmpeg", [self.ffmpeg, "-hide_banner", *([] if feed is not None else ["-nostdin"]), *args], feed is not None)
+        if feed is not None:
+            run.tasks.append(asyncio.ensure_future(_pipe_into(feed, f)))
+        assert f.stdout is not None
+        while chunk := await f.stdout.read(_CHUNK):
+            yield chunk
+
+
+async def _drain(reader: asyncio.StreamReader) -> None:
+    while await reader.read(_CHUNK):
+        pass
+
+
+async def _pipe_into(source: Any, f: asyncio.subprocess.Process) -> None:
+    stdin = f.stdin
+    assert stdin is not None
+    try:
+        async for chunk in _iter_bytes(source):
+            stdin.write(chunk)
+            await stdin.drain()
+    except (ConnectionError, BrokenPipeError, OSError):
+        pass  # ffmpeg went away (skip)
+    finally:
+        try:
+            stdin.close()
+        except (ConnectionError, OSError):
+            pass