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

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 .. 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(Generic[T_co], ObjectReceiveStream[T_co]): 

81 _state: _MemoryObjectStreamState[T_co] 

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

83 

84 def __post_init__(self) -> None: 

85 self._state.open_receive_channels += 1 

86 

87 def receive_nowait(self) -> T_co: 

88 """ 

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

90 

91 :return: the received item 

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

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

94 closed from the sending end 

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

96 waiting to send 

97 

98 """ 

99 if self._closed: 

100 raise ClosedResourceError 

101 

102 if self._state.waiting_senders: 

103 # Get the item from the next sender 

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

105 self._state.buffer.append(item) 

106 send_event.set() 

107 

108 if self._state.buffer: 

109 return self._state.buffer.popleft() 

110 elif not self._state.open_send_channels: 

111 raise EndOfStream 

112 

113 raise WouldBlock 

114 

115 async def receive(self) -> T_co: 

116 await checkpoint() 

117 try: 

118 return self.receive_nowait() 

119 except WouldBlock: 

120 # Add ourselves in the queue 

121 receive_event = Event() 

122 receiver = _MemoryObjectItemReceiver[T_co]() 

123 self._state.waiting_receivers[receive_event] = receiver 

124 

125 try: 

126 await receive_event.wait() 

127 finally: 

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

129 

130 try: 

131 return receiver.item 

132 except AttributeError: 

133 raise EndOfStream from None 

134 

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

136 """ 

137 Create a clone of this receive stream. 

138 

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

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

141 

142 :return: the cloned stream 

143 

144 """ 

145 if self._closed: 

146 raise ClosedResourceError 

147 

148 return MemoryObjectReceiveStream(_state=self._state) 

149 

150 def close(self) -> None: 

151 """ 

152 Close the stream. 

153 

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

155 case for the benefit of synchronous callbacks. 

156 

157 """ 

158 if not self._closed: 

159 self._closed = True 

160 self._state.open_receive_channels -= 1 

161 if self._state.open_receive_channels == 0: 

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

163 for event in send_events: 

164 event.set() 

165 

166 async def aclose(self) -> None: 

167 self.close() 

168 

169 def statistics(self) -> MemoryObjectStreamStatistics: 

170 """ 

171 Return statistics about the current state of this stream. 

172 

173 .. versionadded:: 3.0 

174 """ 

175 return self._state.statistics() 

176 

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

178 return self 

179 

180 def __exit__( 

181 self, 

182 exc_type: type[BaseException] | None, 

183 exc_val: BaseException | None, 

184 exc_tb: TracebackType | None, 

185 ) -> None: 

186 self.close() 

187 

188 def __del__(self) -> None: 

189 if not self._closed: 

190 warnings.warn( 

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

192 ResourceWarning, 

193 stacklevel=1, 

194 source=self, 

195 ) 

196 

197 

198@dataclass(eq=False) 

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

200 _state: _MemoryObjectStreamState[T_contra] 

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

202 

203 def __post_init__(self) -> None: 

204 self._state.open_send_channels += 1 

205 

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

207 """ 

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

209 

210 :param item: the item to send 

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

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

213 receiving end 

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

215 to receive 

216 

217 """ 

218 if self._closed: 

219 raise ClosedResourceError 

220 if not self._state.open_receive_channels: 

221 raise BrokenResourceError 

222 

223 while self._state.waiting_receivers: 

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

225 if not receiver.task_info.has_pending_cancellation(): 

226 receiver.item = item 

227 receive_event.set() 

228 return 

229 

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

231 self._state.buffer.append(item) 

232 else: 

233 raise WouldBlock 

234 

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

236 """ 

237 Send an item to the stream. 

238 

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

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

241 

242 :param item: the item to send 

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

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

245 receiving end 

246 

247 """ 

248 await checkpoint() 

249 try: 

250 self.send_nowait(item) 

251 except WouldBlock: 

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

253 send_event = Event() 

254 self._state.waiting_senders[send_event] = item 

255 try: 

256 await send_event.wait() 

257 except BaseException: 

258 self._state.waiting_senders.pop(send_event, None) 

259 raise 

260 

261 if send_event in self._state.waiting_senders: 

262 del self._state.waiting_senders[send_event] 

263 raise BrokenResourceError from None 

264 

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

266 """ 

267 Create a clone of this send stream. 

268 

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

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

271 

272 :return: the cloned stream 

273 

274 """ 

275 if self._closed: 

276 raise ClosedResourceError 

277 

278 return MemoryObjectSendStream(_state=self._state) 

279 

280 def close(self) -> None: 

281 """ 

282 Close the stream. 

283 

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

285 case for the benefit of synchronous callbacks. 

286 

287 """ 

288 if not self._closed: 

289 self._closed = True 

290 self._state.open_send_channels -= 1 

291 if self._state.open_send_channels == 0: 

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

293 self._state.waiting_receivers.clear() 

294 for event in receive_events: 

295 event.set() 

296 

297 async def aclose(self) -> None: 

298 self.close() 

299 

300 def statistics(self) -> MemoryObjectStreamStatistics: 

301 """ 

302 Return statistics about the current state of this stream. 

303 

304 .. versionadded:: 3.0 

305 """ 

306 return self._state.statistics() 

307 

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

309 return self 

310 

311 def __exit__( 

312 self, 

313 exc_type: type[BaseException] | None, 

314 exc_val: BaseException | None, 

315 exc_tb: TracebackType | None, 

316 ) -> None: 

317 self.close() 

318 

319 def __del__(self) -> None: 

320 if not self._closed: 

321 warnings.warn( 

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

323 ResourceWarning, 

324 stacklevel=1, 

325 source=self, 

326 )