Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/anyio/streams/memory.py: 36%

1from __future__ import annotations 

2 

3__all__ = ( 

4 "MemoryObjectReceiveStream", 

5 "MemoryObjectSendStream", 

6 "MemoryObjectStreamStatistics", 

7) 

8 

9import warnings 

10from collections import OrderedDict, deque 

11from dataclasses import dataclass, field 

12from types import TracebackType 

13from typing import Generic, NamedTuple, TypeVar 

14 

15from .._core._exceptions import ( 

16 BrokenResourceError, 

17 ClosedResourceError, 

18 EndOfStream, 

19 WouldBlock, 

20) 

21from .._core._synchronization import Event 

22from .._core._testing import TaskInfo, get_current_task 

23from ..abc import ObjectReceiveStream, ObjectSendStream 

24from ..lowlevel import checkpoint 

25 

26T_Item = TypeVar("T_Item") 

27T_co = TypeVar("T_co", covariant=True) 

28T_contra = TypeVar("T_contra", contravariant=True) 

29 

30 

31class MemoryObjectStreamStatistics(NamedTuple): 

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

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

34 max_buffer_size: float 

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

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

37 #: number of tasks blocked on :meth:`MemoryObjectSendStream.send` 

38 tasks_waiting_send: int 

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

40 tasks_waiting_receive: int 

41 

42 

43@dataclass(eq=False) 

44class _MemoryObjectItemReceiver(Generic[T_Item]): 

45 task_info: TaskInfo = field(init=False, default_factory=get_current_task) 

46 item: T_Item = field(init=False) 

47 

48 def __repr__(self) -> str: 

49 # When item is not defined, we get following error with default __repr__: 

50 # AttributeError: 'MemoryObjectItemReceiver' object has no attribute 'item' 

51 item = getattr(self, "item", None) 

52 return f"{self.__class__.__name__}(task_info={self.task_info}, item={item!r})" 

53 

54 

55@dataclass(eq=False) 

56class _MemoryObjectStreamState(Generic[T_Item]): 

57 max_buffer_size: float = field() 

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

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

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

61 waiting_receivers: OrderedDict[Event, _MemoryObjectItemReceiver[T_Item]] = field( 

62 init=False, default_factory=OrderedDict 

63 ) 

64 waiting_senders: OrderedDict[Event, T_Item] = field( 

65 init=False, default_factory=OrderedDict 

66 ) 

67 

68 def statistics(self) -> MemoryObjectStreamStatistics: 

69 return MemoryObjectStreamStatistics( 

70 len(self.buffer), 

71 self.max_buffer_size, 

72 self.open_send_channels, 

73 self.open_receive_channels, 

74 len(self.waiting_senders), 

75 len(self.waiting_receivers), 

76 ) 

77 

78 

79@dataclass(eq=False) 

80class MemoryObjectReceiveStream(ObjectReceiveStream[T_co], Generic[T_co]): 

81 _state: _MemoryObjectStreamState[T_co] 

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

83 _waiters: set[Event] | None = field(init=False, default=None) 

84 

85 def __post_init__(self) -> None: 

86 self._state.open_receive_channels += 1 

87 

88 def receive_nowait(self) -> T_co: 

89 """ 

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

91 

92 :return: the received item 

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

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

95 closed from the sending end 

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

97 waiting to send 

98 

99 """ 

100 if self._closed: 

101 raise ClosedResourceError 

102 

103 if self._state.waiting_senders: 

104 # Get the item from the next sender 

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

106 self._state.buffer.append(item) 

107 send_event.set() 

108 

109 if self._state.buffer: 

110 return self._state.buffer.popleft() 

111 elif not self._state.open_send_channels: 

112 raise EndOfStream 

113 

114 raise WouldBlock 

115 

116 async def receive(self) -> T_co: 

117 await checkpoint() 

118 try: 

119 return self.receive_nowait() 

120 except WouldBlock: 

121 if self._waiters is None: 

122 self._waiters = set[Event]() 

123 

124 # Add ourselves in the queue 

125 receive_event = Event() 

126 self._waiters.add(receive_event) 

127 receiver = _MemoryObjectItemReceiver[T_co]() 

128 self._state.waiting_receivers[receive_event] = receiver 

129 

130 try: 

131 await receive_event.wait() 

132 finally: 

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

134 if self._waiters is not None: 

135 self._waiters.remove(receive_event) 

136 

137 try: 

138 return receiver.item 

139 except AttributeError: 

140 if self._closed: 

141 raise ClosedResourceError from None 

142 else: 

143 raise EndOfStream from None 

144 

145 def clone(self) -> MemoryObjectReceiveStream[T_co]: 

146 """ 

147 Create a clone of this receive stream. 

148 

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

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

151 

152 :return: the cloned stream 

153 

154 """ 

155 if self._closed: 

156 raise ClosedResourceError 

157 

158 return MemoryObjectReceiveStream(_state=self._state) 

159 

160 def close(self) -> None: 

161 """ 

162 Close the stream. 

163 

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

165 case for the benefit of synchronous callbacks. 

166 

167 """ 

168 if not self._closed: 

169 self._closed = True 

170 self._state.open_receive_channels -= 1 

171 if self._state.open_receive_channels == 0: 

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

173 for event in send_events: 

174 event.set() 

175 

176 if self._waiters is not None: 

177 for event in self._waiters: 

178 self._state.waiting_receivers.pop(event, None) 

179 event.set() 

180 

181 self._waiters = None 

182 

183 async def aclose(self) -> None: 

184 self.close() 

185 

186 def statistics(self) -> MemoryObjectStreamStatistics: 

187 """ 

188 Return statistics about the current state of this stream. 

189 

190 .. versionadded:: 3.0 

191 """ 

192 return self._state.statistics() 

193 

194 def __enter__(self) -> MemoryObjectReceiveStream[T_co]: 

195 return self 

196 

197 def __exit__( 

198 self, 

199 exc_type: type[BaseException] | None, 

200 exc_val: BaseException | None, 

201 exc_tb: TracebackType | None, 

202 ) -> None: 

203 self.close() 

204 

205 def __del__(self) -> None: 

206 if not self._closed: 

207 warnings.warn( 

208 f"Unclosed <{self.__class__.__name__} at {id(self):x}>", 

209 ResourceWarning, 

210 stacklevel=1, 

211 source=self, 

212 ) 

213 

214 

215@dataclass(eq=False) 

216class MemoryObjectSendStream(ObjectSendStream[T_contra], Generic[T_contra]): 

217 _state: _MemoryObjectStreamState[T_contra] 

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

219 _waiters: set[Event] | None = field(init=False, default=None) 

220 

221 def __post_init__(self) -> None: 

222 self._state.open_send_channels += 1 

223 

224 def send_nowait(self, item: T_contra) -> None: 

225 """ 

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

227 

228 :param item: the item to send 

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

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

231 receiving end 

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

233 to receive 

234 

235 """ 

236 if self._closed: 

237 raise ClosedResourceError 

238 if not self._state.open_receive_channels: 

239 raise BrokenResourceError 

240 

241 while self._state.waiting_receivers: 

242 receive_event, receiver = self._state.waiting_receivers.popitem(last=False) 

243 if not receiver.task_info.has_pending_cancellation(): 

244 receiver.item = item 

245 receive_event.set() 

246 return 

247 

248 if len(self._state.buffer) < self._state.max_buffer_size: 

249 self._state.buffer.append(item) 

250 else: 

251 raise WouldBlock 

252 

253 async def send(self, item: T_contra) -> None: 

254 """ 

255 Send an item to the stream. 

256 

257 If the buffer is full, this method blocks until there is again room in the 

258 buffer or the item can be sent directly to a receiver. 

259 

260 :param item: the item to send 

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

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

263 receiving end 

264 

265 """ 

266 await checkpoint() 

267 try: 

268 self.send_nowait(item) 

269 except WouldBlock: 

270 if self._waiters is None: 

271 self._waiters = set[Event]() 

272 

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

274 send_event = Event() 

275 self._waiters.add(send_event) 

276 self._state.waiting_senders[send_event] = item 

277 try: 

278 await send_event.wait() 

279 finally: 

280 if item_not_consumed := send_event in self._state.waiting_senders: 

281 del self._state.waiting_senders[send_event] 

282 

283 withdrawn = self._closed and send_event in self._waiters 

284 self._waiters.discard(send_event) 

285 

286 if item_not_consumed: 

287 raise BrokenResourceError from None 

288 elif withdrawn: 

289 raise ClosedResourceError from None 

290 

291 def clone(self) -> MemoryObjectSendStream[T_contra]: 

292 """ 

293 Create a clone of this send stream. 

294 

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

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

297 

298 :return: the cloned stream 

299 

300 """ 

301 if self._closed: 

302 raise ClosedResourceError 

303 

304 return MemoryObjectSendStream(_state=self._state) 

305 

306 def close(self) -> None: 

307 """ 

308 Close the stream. 

309 

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

311 case for the benefit of synchronous callbacks. 

312 

313 """ 

314 if not self._closed: 

315 self._closed = True 

316 self._state.open_send_channels -= 1 

317 if self._state.open_send_channels == 0: 

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

319 self._state.waiting_receivers.clear() 

320 for event in receive_events: 

321 event.set() 

322 

323 if self._waiters is not None: 

324 withdrawn = set[Event]() 

325 for event in self._waiters: 

326 if event in self._state.waiting_senders: 

327 del self._state.waiting_senders[event] 

328 withdrawn.add(event) 

329 

330 event.set() 

331 

332 # An event remaining on this set signals that its item was never 

333 # delivered, and send() will raise ClosedResourceError 

334 self._waiters = withdrawn 

335 

336 async def aclose(self) -> None: 

337 self.close() 

338 

339 def statistics(self) -> MemoryObjectStreamStatistics: 

340 """ 

341 Return statistics about the current state of this stream. 

342 

343 .. versionadded:: 3.0 

344 """ 

345 return self._state.statistics() 

346 

347 def __enter__(self) -> MemoryObjectSendStream[T_contra]: 

348 return self 

349 

350 def __exit__( 

351 self, 

352 exc_type: type[BaseException] | None, 

353 exc_val: BaseException | None, 

354 exc_tb: TracebackType | None, 

355 ) -> None: 

356 self.close() 

357 

358 def __del__(self) -> None: 

359 if not self._closed: 

360 warnings.warn( 

361 f"Unclosed <{self.__class__.__name__} at {id(self):x}>", 

362 ResourceWarning, 

363 stacklevel=1, 

364 source=self, 

365 )