208 lines
7.3 KiB
Python
208 lines
7.3 KiB
Python
|
from collections import defaultdict
|
||
|
|
from logging import getLogger
|
||
|
|
from threading import RLock
|
||
|
|
from typing import Dict, Iterable, List, Optional, Union
|
||
|
|
|
||
|
|
from .._model import AudioDevice, DeviceType, StreamType
|
||
|
|
from .._streams import AudioThread
|
||
|
|
from ._device import DeviceManager
|
||
|
|
|
||
|
|
|
||
|
|
class StreamManager:
|
||
|
|
"""
|
||
|
|
The audio manager is responsible for storing the current state of the
|
||
|
|
playing/recording audio streams and allowing fast flexible lookups (by
|
||
|
|
stream index, name, type, device, and any combination of those).
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init__(self, device_manager: DeviceManager):
|
||
|
|
"""
|
||
|
|
:param device_manager: Reference to the device manager.
|
||
|
|
"""
|
||
|
|
self._next_stream_index = 1
|
||
|
|
self._device_manager = device_manager
|
||
|
|
self._state_lock = RLock()
|
||
|
|
self._stream_index_by_name: Dict[str, int] = {}
|
||
|
|
self._stream_name_by_index: Dict[int, str] = {}
|
||
|
|
self._stream_index_to_device: Dict[int, AudioDevice] = {}
|
||
|
|
self._stream_index_to_type: Dict[int, StreamType] = {}
|
||
|
|
self.logger = getLogger(__name__)
|
||
|
|
|
||
|
|
self._streams: Dict[
|
||
|
|
int, Dict[StreamType, Dict[int, AudioThread]]
|
||
|
|
] = defaultdict(lambda: {stream_type: {} for stream_type in StreamType})
|
||
|
|
""" {device_index: {stream_type: {stream_index: audio_thread}}} """
|
||
|
|
|
||
|
|
self._streams_by_index: Dict[StreamType, Dict[int, AudioThread]] = {
|
||
|
|
stream_type: {} for stream_type in StreamType
|
||
|
|
}
|
||
|
|
""" {stream_type: {stream_index: [audio_threads]}} """
|
||
|
|
|
||
|
|
self._stream_locks: Dict[int, Dict[StreamType, RLock]] = defaultdict(
|
||
|
|
lambda: {stream_type: RLock() for stream_type in StreamType}
|
||
|
|
)
|
||
|
|
""" {device_index: {stream_type: RLock}} """
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
def _generate_stream_name(
|
||
|
|
cls,
|
||
|
|
type: StreamType, # pylint: disable=redefined-builtin
|
||
|
|
stream_index: int,
|
||
|
|
) -> str:
|
||
|
|
return f'platypush:audio:{type.value}:{stream_index}'
|
||
|
|
|
||
|
|
def _gen_next_stream_index(
|
||
|
|
self,
|
||
|
|
type: StreamType, # pylint: disable=redefined-builtin
|
||
|
|
stream_name: Optional[str] = None,
|
||
|
|
) -> int:
|
||
|
|
"""
|
||
|
|
:param type: The type of the stream to allocate (input or output).
|
||
|
|
:param stream_name: The name of the stream to allocate.
|
||
|
|
:return: The index of the new stream.
|
||
|
|
"""
|
||
|
|
with self._state_lock:
|
||
|
|
stream_index = self._next_stream_index
|
||
|
|
|
||
|
|
if not stream_name:
|
||
|
|
stream_name = self._generate_stream_name(type, stream_index)
|
||
|
|
|
||
|
|
self._stream_name_by_index[stream_index] = stream_name
|
||
|
|
self._stream_index_by_name[stream_name] = stream_index
|
||
|
|
self._next_stream_index += 1
|
||
|
|
|
||
|
|
return stream_index
|
||
|
|
|
||
|
|
def register(
|
||
|
|
self,
|
||
|
|
audio_thread: AudioThread,
|
||
|
|
device: AudioDevice,
|
||
|
|
type: StreamType, # pylint: disable=redefined-builtin
|
||
|
|
stream_name: Optional[str] = None,
|
||
|
|
):
|
||
|
|
"""
|
||
|
|
Registers an audio stream to a device.
|
||
|
|
|
||
|
|
:param audio_thread: Stream to register.
|
||
|
|
:param device: Device to register the stream to.
|
||
|
|
:param type: The type of the stream to allocate (input or output).
|
||
|
|
:param stream_name: The name of the stream to allocate.
|
||
|
|
"""
|
||
|
|
with self._state_lock:
|
||
|
|
stream_index = audio_thread.stream_index
|
||
|
|
if stream_index is None:
|
||
|
|
stream_index = audio_thread.stream_index = self._gen_next_stream_index(
|
||
|
|
type, stream_name=stream_name
|
||
|
|
)
|
||
|
|
|
||
|
|
self._streams[device.index][type][stream_index] = audio_thread
|
||
|
|
self._stream_index_to_device[stream_index] = device
|
||
|
|
self._stream_index_to_type[stream_index] = type
|
||
|
|
self._streams_by_index[type][stream_index] = audio_thread
|
||
|
|
|
||
|
|
def unregister(
|
||
|
|
self,
|
||
|
|
audio_thread: AudioThread,
|
||
|
|
device: Optional[AudioDevice] = None,
|
||
|
|
type: Optional[StreamType] = None, # pylint: disable=redefined-builtin
|
||
|
|
):
|
||
|
|
"""
|
||
|
|
Unregisters an audio stream from a device.
|
||
|
|
|
||
|
|
:param audio_thread: Stream to unregister.
|
||
|
|
:param device: Device to unregister the stream from.
|
||
|
|
:param type: The type of the stream to unregister (input or output).
|
||
|
|
"""
|
||
|
|
with self._state_lock:
|
||
|
|
stream_index = audio_thread.stream_index
|
||
|
|
if stream_index is None:
|
||
|
|
return
|
||
|
|
|
||
|
|
if device is None:
|
||
|
|
device = self._stream_index_to_device.get(stream_index)
|
||
|
|
|
||
|
|
if not type:
|
||
|
|
type = self._stream_index_to_type.get(stream_index)
|
||
|
|
|
||
|
|
if device is None or type is None:
|
||
|
|
return
|
||
|
|
|
||
|
|
self._streams[device.index][type].pop(stream_index, None)
|
||
|
|
self._stream_index_to_device.pop(stream_index, None)
|
||
|
|
self._stream_index_to_type.pop(stream_index, None)
|
||
|
|
self._streams_by_index[type].pop(stream_index, None)
|
||
|
|
stream_name = self._stream_name_by_index.pop(stream_index, None)
|
||
|
|
if stream_name:
|
||
|
|
self._stream_index_by_name.pop(stream_name, None)
|
||
|
|
|
||
|
|
def _get_by_device_and_type(
|
||
|
|
self,
|
||
|
|
device: Optional[DeviceType] = None,
|
||
|
|
type: Optional[StreamType] = None, # pylint: disable=redefined-builtin
|
||
|
|
) -> List[AudioThread]:
|
||
|
|
"""
|
||
|
|
Filter streams by device and/or type.
|
||
|
|
"""
|
||
|
|
devs = (
|
||
|
|
[self._device_manager.get_device(device, type)]
|
||
|
|
if device is not None
|
||
|
|
else self._device_manager.get_devices(type)
|
||
|
|
)
|
||
|
|
|
||
|
|
return [
|
||
|
|
audio_thread
|
||
|
|
for dev in devs
|
||
|
|
for stream_info in (
|
||
|
|
[self._streams[dev.index].get(type, {})]
|
||
|
|
if type
|
||
|
|
else list(self._streams[dev.index].values())
|
||
|
|
)
|
||
|
|
for audio_thread in stream_info.values()
|
||
|
|
if audio_thread and audio_thread.is_alive()
|
||
|
|
]
|
||
|
|
|
||
|
|
def _get_by_stream_index_or_name(
|
||
|
|
self, streams: Iterable[Union[str, int]]
|
||
|
|
) -> List[AudioThread]:
|
||
|
|
"""
|
||
|
|
Filter streams by index or name.
|
||
|
|
"""
|
||
|
|
threads = []
|
||
|
|
|
||
|
|
for stream in streams:
|
||
|
|
try:
|
||
|
|
stream_index = int(stream)
|
||
|
|
except (TypeError, ValueError):
|
||
|
|
stream_index = self._stream_index_by_name.get(stream) # type: ignore
|
||
|
|
if stream_index is None:
|
||
|
|
self.logger.warning('No such audio stream: %s', stream)
|
||
|
|
continue
|
||
|
|
|
||
|
|
stream_type = self._stream_index_to_type.get(stream_index)
|
||
|
|
if not stream_type:
|
||
|
|
self.logger.warning(
|
||
|
|
'No type available for this audio stream: %s', stream
|
||
|
|
)
|
||
|
|
continue
|
||
|
|
|
||
|
|
thread = self._streams_by_index.get(stream_type, {}).get(stream_index)
|
||
|
|
if thread:
|
||
|
|
threads.append(thread)
|
||
|
|
|
||
|
|
return threads
|
||
|
|
|
||
|
|
def get(
|
||
|
|
self,
|
||
|
|
device: Optional[DeviceType] = None,
|
||
|
|
type: Optional[StreamType] = None, # pylint: disable=redefined-builtin
|
||
|
|
streams: Optional[Iterable[Union[str, int]]] = None,
|
||
|
|
) -> List[AudioThread]:
|
||
|
|
"""
|
||
|
|
Searches streams, either by device and/or type, or by stream index/name.
|
||
|
|
"""
|
||
|
|
return (
|
||
|
|
self._get_by_stream_index_or_name(streams)
|
||
|
|
if streams
|
||
|
|
else self._get_by_device_and_type(device, type)
|
||
|
|
)
|