332 lines
9.4 KiB
Python
332 lines
9.4 KiB
Python
from abc import ABC, abstractmethod
|
|
import asyncio
|
|
from asyncio.subprocess import PIPE
|
|
from logging import getLogger
|
|
from queue import Empty, Queue
|
|
from threading import Event, RLock, Thread
|
|
from typing import Any, Callable, Coroutine, Iterable, Optional, Self
|
|
|
|
from platypush.context import get_or_create_event_loop
|
|
|
|
_dtype_to_ffmpeg_format = {
|
|
'int8': 's8',
|
|
'uint8': 'u8',
|
|
'int16': 's16le',
|
|
'uint16': 'u16le',
|
|
'int32': 's32le',
|
|
'uint32': 'u32le',
|
|
'float32': 'f32le',
|
|
'float64': 'f64le',
|
|
}
|
|
"""
|
|
Supported raw types:
|
|
'int8', 'uint8', 'int16', 'uint16', 'int32', 'uint32', 'float32', 'float64'
|
|
"""
|
|
|
|
|
|
class AudioConverter(Thread, ABC):
|
|
"""
|
|
Base class for an ffmpeg audio converter instance.
|
|
"""
|
|
|
|
_format_to_ffmpeg_args = {
|
|
'wav': ('-f', 'wav'),
|
|
'ogg': ('-f', 'ogg'),
|
|
'mp3': ('-f', 'mp3'),
|
|
'aac': ('-f', 'adts'),
|
|
'flac': ('-f', 'flac'),
|
|
}
|
|
|
|
def __init__(
|
|
self,
|
|
*args,
|
|
ffmpeg_bin: str,
|
|
sample_rate: int,
|
|
channels: int,
|
|
volume: float,
|
|
dtype: str,
|
|
chunk_size: int,
|
|
format: Optional[str] = None, # pylint: disable=redefined-builtin
|
|
on_exit: Optional[Callable[[], Any]] = None,
|
|
**kwargs,
|
|
):
|
|
"""
|
|
:param ffmpeg_bin: Path to the ffmpeg binary.
|
|
:param sample_rate: The sample rate of the input/output audio.
|
|
:param channels: The number of channels of the input/output audio.
|
|
:param volume: Audio volume, as a percentage between 0 and 100.
|
|
:param dtype: The (numpy) data type of the raw input/output audio.
|
|
:param chunk_size: Number of bytes that will be read at once from the
|
|
ffmpeg process.
|
|
:param format: Input/output audio format.
|
|
:param on_exit: Function to call when the ffmpeg process exits.
|
|
"""
|
|
super().__init__(*args, **kwargs)
|
|
|
|
ffmpeg_format = _dtype_to_ffmpeg_format.get(dtype)
|
|
assert ffmpeg_format, (
|
|
f'Unsupported data type: {dtype}. Supported data types: '
|
|
f'{list(_dtype_to_ffmpeg_format.keys())}'
|
|
)
|
|
|
|
self._ffmpeg_bin = ffmpeg_bin
|
|
self._ffmpeg_format = ffmpeg_format
|
|
self._ffmpeg_task: Optional[Coroutine] = None
|
|
self._sample_rate = sample_rate
|
|
self._channels = channels
|
|
self._chunk_size = chunk_size
|
|
self._format = format
|
|
self._closed = False
|
|
self._out_queue = Queue()
|
|
self.ffmpeg = None
|
|
self.volume = volume
|
|
self.logger = getLogger(__name__)
|
|
self._loop = None
|
|
self._should_stop = Event()
|
|
self._stop_lock = RLock()
|
|
self._on_exit = on_exit
|
|
self._ffmpeg_terminated = Event()
|
|
|
|
def __enter__(self) -> Self:
|
|
"""
|
|
Audio converter context manager.
|
|
|
|
It starts and registers the ffmpeg converter process.
|
|
"""
|
|
self.start()
|
|
return self
|
|
|
|
def __exit__(self, *_, **__):
|
|
"""
|
|
Audio converter context manager.
|
|
|
|
It stops and unregisters the ffmpeg converter process.
|
|
"""
|
|
self.stop()
|
|
|
|
def _check_ffmpeg(self):
|
|
assert not self.terminated, 'The ffmpeg process has already terminated'
|
|
|
|
@property
|
|
def gain(self) -> float:
|
|
return self.volume / 100
|
|
|
|
@property
|
|
def terminated(self) -> bool:
|
|
return self._ffmpeg_terminated.is_set()
|
|
|
|
@property
|
|
def _default_args(self) -> Iterable[str]:
|
|
"""
|
|
Set of arguments common to all ffmpeg converter instances.
|
|
"""
|
|
return ('-hide_banner', '-loglevel', 'warning', '-y')
|
|
|
|
@property
|
|
@abstractmethod
|
|
def _input_format_args(self) -> Iterable[str]:
|
|
"""
|
|
Ffmpeg audio input arguments.
|
|
"""
|
|
raise NotImplementedError()
|
|
|
|
@property
|
|
@abstractmethod
|
|
def _output_format_args(self):
|
|
"""
|
|
Ffmpeg audio output arguments.
|
|
"""
|
|
raise NotImplementedError()
|
|
|
|
@property
|
|
def _channel_layout_args(self) -> Iterable[str]:
|
|
"""
|
|
Set of extra ffmpeg arguments for the channel layout.
|
|
"""
|
|
args = ('-ac', str(self._channels))
|
|
if self._channels == 1:
|
|
return args + ('-channel_layout', 'mono')
|
|
if self._channels == 2:
|
|
return args + ('-channel_layout', 'stereo')
|
|
return args
|
|
|
|
@property
|
|
def _raw_ffmpeg_args(self) -> Iterable[str]:
|
|
"""
|
|
Ffmpeg arguments for raw audio input/output given the current
|
|
configuration.
|
|
"""
|
|
return (
|
|
'-f',
|
|
self._ffmpeg_format,
|
|
'-ar',
|
|
str(self._sample_rate),
|
|
*self._channel_layout_args,
|
|
)
|
|
|
|
@property
|
|
def _audio_volume_args(self) -> Iterable[str]:
|
|
"""
|
|
Ffmpeg audio volume arguments.
|
|
"""
|
|
return ('-filter:a', f'volume={self.gain}')
|
|
|
|
@property
|
|
def _input_source_args(self) -> Iterable[str]:
|
|
"""
|
|
Default arguments for the ffmpeg input source (default: ``-i pipe:``,
|
|
ffmpeg will read from a pipe filled by the application).
|
|
"""
|
|
return ('-i', 'pipe:')
|
|
|
|
@property
|
|
def _output_target_args(self) -> Iterable[str]:
|
|
"""
|
|
Default arguments for the ffmpeg output target (default: ``pipe:``,
|
|
ffmpeg will write the output to a pipe read by the application).
|
|
"""
|
|
return ('pipe:',)
|
|
|
|
@property
|
|
def _converter_stdin(self) -> Optional[int]:
|
|
"""
|
|
Default stdin file descriptor to be used by the ffmpeg converter.
|
|
|
|
Default: ``PIPE``, as the ffmpeg process by default reads audio frames
|
|
from the stdin.
|
|
"""
|
|
return PIPE
|
|
|
|
@property
|
|
def _compressed_ffmpeg_args(self) -> Iterable[str]:
|
|
"""
|
|
Ffmpeg arguments for the compressed audio given the current
|
|
configuration.
|
|
"""
|
|
if not self._format:
|
|
return ()
|
|
|
|
ffmpeg_args = self._format_to_ffmpeg_args.get(self._format)
|
|
assert ffmpeg_args, (
|
|
f'Unsupported output format: {self._format}. Supported formats: '
|
|
f'{list(self._format_to_ffmpeg_args.keys())}'
|
|
)
|
|
|
|
return ffmpeg_args
|
|
|
|
async def _audio_proxy(self, timeout: Optional[float] = None):
|
|
"""
|
|
Proxy the converted audio stream to the output queue for downstream
|
|
consumption.
|
|
"""
|
|
ffmpeg_args = (
|
|
self._ffmpeg_bin,
|
|
*self._default_args,
|
|
*self._input_format_args,
|
|
*self._input_source_args,
|
|
*self._output_format_args,
|
|
*self._output_target_args,
|
|
)
|
|
|
|
self.ffmpeg = await asyncio.create_subprocess_exec(
|
|
*ffmpeg_args,
|
|
stdin=self._converter_stdin,
|
|
stdout=PIPE,
|
|
)
|
|
|
|
self.logger.info('Running ffmpeg: %s', ' '.join(ffmpeg_args))
|
|
|
|
try:
|
|
await asyncio.wait_for(self.ffmpeg.wait(), 0.1)
|
|
except asyncio.TimeoutError:
|
|
pass
|
|
|
|
while (
|
|
self._loop
|
|
and self.ffmpeg
|
|
and self.ffmpeg.returncode is None
|
|
and not self.should_stop
|
|
):
|
|
self._check_ffmpeg()
|
|
assert (
|
|
self.ffmpeg and self.ffmpeg.stdout
|
|
), 'The stdout is closed for the ffmpeg process'
|
|
|
|
self._ffmpeg_terminated.clear()
|
|
try:
|
|
data = await asyncio.wait_for(
|
|
self.ffmpeg.stdout.read(self._chunk_size), timeout
|
|
)
|
|
self._out_queue.put(data)
|
|
except asyncio.TimeoutError:
|
|
self._out_queue.put(b'')
|
|
|
|
def write(self, data: bytes):
|
|
"""
|
|
Write raw data to the ffmpeg process.
|
|
"""
|
|
self._check_ffmpeg()
|
|
assert (
|
|
self.ffmpeg and self._loop and self.ffmpeg.stdin
|
|
), 'The stdin is closed for the ffmpeg process'
|
|
|
|
self._loop.call_soon_threadsafe(self.ffmpeg.stdin.write, data)
|
|
|
|
def read(self, timeout: Optional[float] = None) -> Optional[bytes]:
|
|
"""
|
|
Read the next chunk of converted audio bytes from the converter queue.
|
|
"""
|
|
try:
|
|
return self._out_queue.get(timeout=timeout)
|
|
except Empty:
|
|
return None
|
|
|
|
def run(self):
|
|
"""
|
|
Main runner. It runs the audio proxy in a loop and cleans up everything
|
|
in case of stop/failure.
|
|
"""
|
|
super().run()
|
|
self._loop = get_or_create_event_loop()
|
|
try:
|
|
self._ffmpeg_task = self._audio_proxy(timeout=1)
|
|
self._loop.run_until_complete(self._ffmpeg_task)
|
|
except RuntimeError as e:
|
|
self.logger.warning(e)
|
|
finally:
|
|
self.stop()
|
|
|
|
def stop(self):
|
|
"""
|
|
Sets the stop event, kills the ffmpeg process and resets the context.
|
|
"""
|
|
with self._stop_lock:
|
|
self._should_stop.set()
|
|
if self._ffmpeg_task:
|
|
self._ffmpeg_task.close()
|
|
self._ffmpeg_task = None
|
|
|
|
try:
|
|
if self.ffmpeg and self.ffmpeg.returncode is None:
|
|
self.ffmpeg.kill()
|
|
except ProcessLookupError:
|
|
pass
|
|
|
|
self.ffmpeg = None
|
|
self._loop = None
|
|
|
|
self._ffmpeg_terminated.set()
|
|
|
|
if self._on_exit:
|
|
self._on_exit()
|
|
|
|
@property
|
|
def should_stop(self) -> bool:
|
|
"""
|
|
Proxy property for the ``_should_stop`` event.
|
|
"""
|
|
return self._should_stop.is_set()
|
|
|
|
|
|
# vim:sw=4:ts=4:et:
|