1 from collections import OrderedDict, deque

2 from dataclasses import dataclass, field

3 from types import TracebackType

4 from typing import Deque, Generic, List, NamedTuple, Optional, Type, TypeVar

5

6 from .. import (

7 BrokenResourceError,

8 ClosedResourceError,

9 EndOfStream,

10 WouldBlock,

11 get_cancelled_exc_class,

12)

13 from .._core._compat import DeprecatedAwaitable

14 from ..abc import Event, ObjectReceiveStream, ObjectSendStream

15 from ..lowlevel import checkpoint

16

17 T_Item = TypeVar("T_Item")

18

19

20 class MemoryObjectStreamStatistics(NamedTuple):

21 current_buffer_used: int #: number of items stored in the buffer

22 #: maximum number of items that can be stored on this stream (or :data:`math.inf`)

23 max_buffer_size: float

24 open_send_streams: int #: number of unclosed clones of the send stream

25 open_receive_streams: int #: number of unclosed clones of the receive stream

26 tasks_waiting_send: int #: number of tasks blocked on :meth:`MemoryObjectSendStream.send`

27 #: number of tasks blocked on :meth:`MemoryObjectReceiveStream.receive`

28 tasks_waiting_receive: int

29

30

31@dataclass(eq=False)

32 class MemoryObjectStreamState(Generic[T_Item]):

33 max_buffer_size: float = field()

34 buffer: Deque[T_Item] = field(init=False, default_factory=deque)

35 open_send_channels: int = field(init=False, default=0)

36 open_receive_channels: int = field(init=False, default=0)

37 waiting_receivers: "OrderedDict[Event, List[T_Item]]" = field(

38 init=False, default_factory=OrderedDict

39 )

40 waiting_senders: "OrderedDict[Event, T_Item]" = field(

41 init=False, default_factory=OrderedDict

42 )

43

44 def statistics(self) -> MemoryObjectStreamStatistics:

45 return MemoryObjectStreamStatistics(

46 len(self.buffer),

47 self.max_buffer_size,

48 self.open_send_channels,

49 self.open_receive_channels,

50 len(self.waiting_senders),

51 len(self.waiting_receivers),

52 )

53

54

55@dataclass(eq=False)

56 class MemoryObjectReceiveStream(Generic[T_Item], ObjectReceiveStream[T_Item]):

57 _state: MemoryObjectStreamState[T_Item]

58 _closed: bool = field(init=False, default=False)

59

60 def __post_init__(self) -> None:

61 self._state.open_receive_channels += 1

62

63 def receive_nowait(self) -> T_Item:

64 """

65 Receive the next item if it can be done without waiting.

66

67 :return: the received item

68 :raises ~anyio.ClosedResourceError: if this send stream has been closed

69 :raises ~anyio.EndOfStream: if the buffer is empty and this stream has been

70 closed from the sending end

71 :raises ~anyio.WouldBlock: if there are no items in the buffer and no tasks

72 waiting to send

73

74 """

75 if self._closed:

76 raise ClosedResourceError

77

78 if self._state.waiting_senders:

79 # Get the item from the next sender

80 send_event, item = self._state.waiting_senders.popitem(last=False)

81 self._state.buffer.append(item)

82 send_event.set()

83

84 if self._state.buffer:

85 return self._state.buffer.popleft()

86 elif not self._state.open_send_channels:

87 raise EndOfStream

88

89 raise WouldBlock

90

91 async def receive(self) -> T_Item:

92 await checkpoint()

93 try:

94 return self.receive_nowait()

95 except WouldBlock:

96 # Add ourselves in the queue

97 receive_event = Event()

98 container: List[T_Item] = []

99 self._state.waiting_receivers[receive_event] = container

100

101 try:

102 await receive_event.wait()

103 except get_cancelled_exc_class():

104 # Ignore the immediate cancellation if we already received an item, so as not to

105 # lose it

106 if not container:

107 raise

108 finally:

109 self._state.waiting_receivers.pop(receive_event, None)

110

111 if container:

112 return container[0]

113 else:

114 raise EndOfStream

115

116 def clone(self) -> "MemoryObjectReceiveStream[T_Item]":

117 """

118 Create a clone of this receive stream.

119

120 Each clone can be closed separately. Only when all clones have been closed will the

121 receiving end of the memory stream be considered closed by the sending ends.

122

123 :return: the cloned stream

124

125 """

126 if self._closed:

127 raise ClosedResourceError

128

129 return MemoryObjectReceiveStream(_state=self._state)

130

131 def close(self) -> None:

132 """

133 Close the stream.

134

135 This works the exact same way as :meth:`aclose`, but is provided as a special case for the

136 benefit of synchronous callbacks.

137

138 """

139 if not self._closed:

140 self._closed = True

141 self._state.open_receive_channels -= 1

142 if self._state.open_receive_channels == 0:

143 send_events = list(self._state.waiting_senders.keys())

144 for event in send_events:

145 event.set()

146

147 async def aclose(self) -> None:

148 self.close()

149

150 def statistics(self) -> MemoryObjectStreamStatistics:

151 """

152 Return statistics about the current state of this stream.

153

154 .. versionadded:: 3.0

155 """

156 return self._state.statistics()

157

158 def __enter__(self) -> "MemoryObjectReceiveStream[T_Item]":

159 return self

160

161 def __exit__(

162 self,

163 exc_type: Optional[Type[BaseException]],

164 exc_val: Optional[BaseException],

165 exc_tb: Optional[TracebackType],

166 ) -> None:

167 self.close()

168

169

170@dataclass(eq=False)

171 class MemoryObjectSendStream(Generic[T_Item], ObjectSendStream[T_Item]):

172 _state: MemoryObjectStreamState[T_Item]

173 _closed: bool = field(init=False, default=False)

174

175 def __post_init__(self) -> None:

176 self._state.open_send_channels += 1

177

178 def send_nowait(self, item: T_Item) -> DeprecatedAwaitable:

179 """

180 Send an item immediately if it can be done without waiting.

181

182 :param item: the item to send

183 :raises ~anyio.ClosedResourceError: if this send stream has been closed

184 :raises ~anyio.BrokenResourceError: if the stream has been closed from the

185 receiving end

186 :raises ~anyio.WouldBlock: if the buffer is full and there are no tasks waiting

187 to receive

188

189 """

190 if self._closed:

191 raise ClosedResourceError

192 if not self._state.open_receive_channels:

193 raise BrokenResourceError

194

195 if self._state.waiting_receivers:

196 receive_event, container = self._state.waiting_receivers.popitem(last=False)

197 container.append(item)

198 receive_event.set()

199 elif len(self._state.buffer) < self._state.max_buffer_size:

200 self._state.buffer.append(item)

201 else:

202 raise WouldBlock

203

204 return DeprecatedAwaitable(self.send_nowait)

205

206 async def send(self, item: T_Item) -> None:

207 await checkpoint()

208 try:

209 self.send_nowait(item)

210 except WouldBlock:

211 # Wait until there's someone on the receiving end

212 send_event = Event()

213 self._state.waiting_senders[send_event] = item

214 try:

215 await send_event.wait()

216 except BaseException:

217 self._state.waiting_senders.pop(send_event, None) # type: ignore[arg-type]

218 raise

219

220 if self._state.waiting_senders.pop(send_event, None): # type: ignore[arg-type]

221 raise BrokenResourceError

222

223 def clone(self) -> "MemoryObjectSendStream[T_Item]":

224 """

225 Create a clone of this send stream.

226

227 Each clone can be closed separately. Only when all clones have been closed will the

228 sending end of the memory stream be considered closed by the receiving ends.

229

230 :return: the cloned stream

231

232 """

233 if self._closed:

234 raise ClosedResourceError

235

236 return MemoryObjectSendStream(_state=self._state)

237

238 def close(self) -> None:

239 """

240 Close the stream.

241

242 This works the exact same way as :meth:`aclose`, but is provided as a special case for the

243 benefit of synchronous callbacks.

244

245 """

246 if not self._closed:

247 self._closed = True

248 self._state.open_send_channels -= 1

249 if self._state.open_send_channels == 0:

250 receive_events = list(self._state.waiting_receivers.keys())

251 self._state.waiting_receivers.clear()

252 for event in receive_events:

253 event.set()

254

255 async def aclose(self) -> None:

256 self.close()

257

258 def statistics(self) -> MemoryObjectStreamStatistics:

259 """

260 Return statistics about the current state of this stream.

261

262 .. versionadded:: 3.0

263 """

264 return self._state.statistics()

265

266 def __enter__(self) -> "MemoryObjectSendStream[T_Item]":

267 return self

268

269 def __exit__(

270 self,

271 exc_type: Optional[Type[BaseException]],

272 exc_val: Optional[BaseException],

273 exc_tb: Optional[TracebackType],

274 ) -> None:

275 self.close()