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) )