from collections import OrderedDict, deque from dataclasses import dataclass, field from types import TracebackType from typing import Deque, Generic, List, NamedTuple, Optional, Type, TypeVar from .. import ( BrokenResourceError, ClosedResourceError, EndOfStream, WouldBlock, get_cancelled_exc_class, ) from .._core._compat import DeprecatedAwaitable from ..abc import Event, ObjectReceiveStream, ObjectSendStream from ..lowlevel import checkpoint T_Item = TypeVar("T_Item") class MemoryObjectStreamStatistics(NamedTuple): current_buffer_used: int #: number of items stored in the buffer #: maximum number of items that can be stored on this stream (or :data:`math.inf`) max_buffer_size: float open_send_streams: int #: number of unclosed clones of the send stream open_receive_streams: int #: number of unclosed clones of the receive stream tasks_waiting_send: int #: number of tasks blocked on :meth:`MemoryObjectSendStream.send` #: number of tasks blocked on :meth:`MemoryObjectReceiveStream.receive` tasks_waiting_receive: int @dataclass(eq=False) class MemoryObjectStreamState(Generic[T_Item]): max_buffer_size: float = field() buffer: Deque[T_Item] = field(init=False, default_factory=deque) open_send_channels: int = field(init=False, default=0) open_receive_channels: int = field(init=False, default=0) waiting_receivers: "OrderedDict[Event, List[T_Item]]" = field( init=False, default_factory=OrderedDict ) waiting_senders: "OrderedDict[Event, T_Item]" = field( init=False, default_factory=OrderedDict ) def statistics(self) -> MemoryObjectStreamStatistics: return MemoryObjectStreamStatistics( len(self.buffer), self.max_buffer_size, self.open_send_channels, self.open_receive_channels, len(self.waiting_senders), len(self.waiting_receivers), ) @dataclass(eq=False) class MemoryObjectReceiveStream(Generic[T_Item], ObjectReceiveStream[T_Item]): _state: MemoryObjectStreamState[T_Item] _closed: bool = field(init=False, default=False) def __post_init__(self) -> None: self._state.open_receive_channels += 1 def receive_nowait(self) -> T_Item: """ Receive the next item if it can be done without waiting. :return: the received item :raises ~anyio.ClosedResourceError: if this send stream has been closed :raises ~anyio.EndOfStream: if the buffer is empty and this stream has been closed from the sending end :raises ~anyio.WouldBlock: if there are no items in the buffer and no tasks waiting to send """ if self._closed: raise ClosedResourceError if self._state.waiting_senders: # Get the item from the next sender send_event, item = self._state.waiting_senders.popitem(last=False) self._state.buffer.append(item) send_event.set() if self._state.buffer: return self._state.buffer.popleft() elif not self._state.open_send_channels: raise EndOfStream raise WouldBlock async def receive(self) -> T_Item: await checkpoint() try: return self.receive_nowait() except WouldBlock: # Add ourselves in the queue receive_event = Event() container: List[T_Item] = [] self._state.waiting_receivers[receive_event] = container try: await receive_event.wait() except get_cancelled_exc_class(): # Ignore the immediate cancellation if we already received an item, so as not to # lose it if not container: raise finally: self._state.waiting_receivers.pop(receive_event, None) if container: return container[0] else: raise EndOfStream def clone(self) -> "MemoryObjectReceiveStream[T_Item]": """ Create a clone of this receive stream. Each clone can be closed separately. Only when all clones have been closed will the receiving end of the memory stream be considered closed by the sending ends. :return: the cloned stream """ if self._closed: raise ClosedResourceError return MemoryObjectReceiveStream(_state=self._state) def close(self) -> None: """ Close the stream. This works the exact same way as :meth:`aclose`, but is provided as a special case for the benefit of synchronous callbacks. """ if not self._closed: self._closed = True self._state.open_receive_channels -= 1 if self._state.open_receive_channels == 0: send_events = list(self._state.waiting_senders.keys()) for event in send_events: event.set() async def aclose(self) -> None: self.close() def statistics(self) -> MemoryObjectStreamStatistics: """ Return statistics about the current state of this stream. .. versionadded:: 3.0 """ return self._state.statistics() def __enter__(self) -> "MemoryObjectReceiveStream[T_Item]": return self def __exit__( self, exc_type: Optional[Type[BaseException]], exc_val: Optional[BaseException], exc_tb: Optional[TracebackType], ) -> None: self.close() @dataclass(eq=False) class MemoryObjectSendStream(Generic[T_Item], ObjectSendStream[T_Item]): _state: MemoryObjectStreamState[T_Item] _closed: bool = field(init=False, default=False) def __post_init__(self) -> None: self._state.open_send_channels += 1 def send_nowait(self, item: T_Item) -> DeprecatedAwaitable: """ Send an item immediately if it can be done without waiting. :param item: the item to send :raises ~anyio.ClosedResourceError: if this send stream has been closed :raises ~anyio.BrokenResourceError: if the stream has been closed from the receiving end :raises ~anyio.WouldBlock: if the buffer is full and there are no tasks waiting to receive """ if self._closed: raise ClosedResourceError if not self._state.open_receive_channels: raise BrokenResourceError if self._state.waiting_receivers: receive_event, container = self._state.waiting_receivers.popitem(last=False) container.append(item) receive_event.set() elif len(self._state.buffer) < self._state.max_buffer_size: self._state.buffer.append(item) else: raise WouldBlock return DeprecatedAwaitable(self.send_nowait) async def send(self, item: T_Item) -> None: await checkpoint() try: self.send_nowait(item) except WouldBlock: # Wait until there's someone on the receiving end send_event = Event() self._state.waiting_senders[send_event] = item try: await send_event.wait() except BaseException: self._state.waiting_senders.pop(send_event, None) # type: ignore[arg-type] raise if self._state.waiting_senders.pop(send_event, None): # type: ignore[arg-type] raise BrokenResourceError def clone(self) -> "MemoryObjectSendStream[T_Item]": """ Create a clone of this send stream. Each clone can be closed separately. Only when all clones have been closed will the sending end of the memory stream be considered closed by the receiving ends. :return: the cloned stream """ if self._closed: raise ClosedResourceError return MemoryObjectSendStream(_state=self._state) def close(self) -> None: """ Close the stream. This works the exact same way as :meth:`aclose`, but is provided as a special case for the benefit of synchronous callbacks. """ if not self._closed: self._closed = True self._state.open_send_channels -= 1 if self._state.open_send_channels == 0: receive_events = list(self._state.waiting_receivers.keys()) self._state.waiting_receivers.clear() for event in receive_events: event.set() async def aclose(self) -> None: self.close() def statistics(self) -> MemoryObjectStreamStatistics: """ Return statistics about the current state of this stream. .. versionadded:: 3.0 """ return self._state.statistics() def __enter__(self) -> "MemoryObjectSendStream[T_Item]": return self def __exit__( self, exc_type: Optional[Type[BaseException]], exc_val: Optional[BaseException], exc_tb: Optional[TracebackType], ) -> None: self.close()