Coverage for /pythoncovmergedfiles/medio/medio/src/aiohttp/aiohttp/client_ws.py: 22%

1"""WebSocket client for asyncio.""" 

2 

3import asyncio 

4import dataclasses 

5import sys 

6from typing import Any, Final, Optional, cast 

7 

8from .client_exceptions import ClientError 

9from .client_reqrep import ClientResponse 

10from .helpers import call_later, set_result 

11from .http import ( 

12 WS_CLOSED_MESSAGE, 

13 WS_CLOSING_MESSAGE, 

14 WebSocketError, 

15 WSCloseCode, 

16 WSMessage, 

17 WSMsgType, 

18) 

19from .http_websocket import WebSocketWriter # WSMessage 

20from .streams import EofStream, FlowControlDataQueue 

21from .typedefs import ( 

22 DEFAULT_JSON_DECODER, 

23 DEFAULT_JSON_ENCODER, 

24 JSONDecoder, 

25 JSONEncoder, 

26) 

27 

28if sys.version_info >= (3, 11): 

29 import asyncio as async_timeout 

30else: 

31 import async_timeout 

32 

33 

34@dataclasses.dataclass(frozen=True) 

35class ClientWSTimeout: 

36 ws_receive: Optional[float] = None 

37 ws_close: Optional[float] = None 

38 

39 

40DEFAULT_WS_CLIENT_TIMEOUT: Final[ClientWSTimeout] = ClientWSTimeout( 

41 ws_receive=None, ws_close=10.0 

42) 

43 

44 

45class ClientWebSocketResponse: 

46 def __init__( 

47 self, 

48 reader: "FlowControlDataQueue[WSMessage]", 

49 writer: WebSocketWriter, 

50 protocol: Optional[str], 

51 response: ClientResponse, 

52 timeout: ClientWSTimeout, 

53 autoclose: bool, 

54 autoping: bool, 

55 loop: asyncio.AbstractEventLoop, 

56 *, 

57 heartbeat: Optional[float] = None, 

58 compress: int = 0, 

59 client_notakeover: bool = False, 

60 ) -> None: 

61 self._response = response 

62 self._conn = response.connection 

63 

64 self._writer = writer 

65 self._reader = reader 

66 self._protocol = protocol 

67 self._closed = False 

68 self._closing = False 

69 self._close_code: Optional[int] = None 

70 self._timeout: ClientWSTimeout = timeout 

71 self._autoclose = autoclose 

72 self._autoping = autoping 

73 self._heartbeat = heartbeat 

74 self._heartbeat_cb: Optional[asyncio.TimerHandle] = None 

75 if heartbeat is not None: 

76 self._pong_heartbeat = heartbeat / 2.0 

77 self._pong_response_cb: Optional[asyncio.TimerHandle] = None 

78 self._loop = loop 

79 self._waiting: Optional[asyncio.Future[bool]] = None 

80 self._exception: Optional[BaseException] = None 

81 self._compress = compress 

82 self._client_notakeover = client_notakeover 

83 

84 self._reset_heartbeat() 

85 

86 def _cancel_heartbeat(self) -> None: 

87 if self._pong_response_cb is not None: 

88 self._pong_response_cb.cancel() 

89 self._pong_response_cb = None 

90 

91 if self._heartbeat_cb is not None: 

92 self._heartbeat_cb.cancel() 

93 self._heartbeat_cb = None 

94 

95 def _reset_heartbeat(self) -> None: 

96 self._cancel_heartbeat() 

97 

98 if self._heartbeat is not None: 

99 self._heartbeat_cb = call_later( 

100 self._send_heartbeat, 

101 self._heartbeat, 

102 self._loop, 

103 timeout_ceil_threshold=self._conn._connector._timeout_ceil_threshold 

104 if self._conn is not None 

105 else 5, 

106 ) 

107 

108 def _send_heartbeat(self) -> None: 

109 if self._heartbeat is not None and not self._closed: 

110 # fire-and-forget a task is not perfect but maybe ok for 

111 # sending ping. Otherwise we need a long-living heartbeat 

112 # task in the class. 

113 self._loop.create_task(self._writer.ping()) # type: ignore[unused-awaitable] 

114 

115 if self._pong_response_cb is not None: 

116 self._pong_response_cb.cancel() 

117 self._pong_response_cb = call_later( 

118 self._pong_not_received, 

119 self._pong_heartbeat, 

120 self._loop, 

121 timeout_ceil_threshold=self._conn._connector._timeout_ceil_threshold 

122 if self._conn is not None 

123 else 5, 

124 ) 

125 

126 def _pong_not_received(self) -> None: 

127 if not self._closed: 

128 self._closed = True 

129 self._close_code = WSCloseCode.ABNORMAL_CLOSURE 

130 self._exception = asyncio.TimeoutError() 

131 self._response.close() 

132 

133 @property 

134 def closed(self) -> bool: 

135 return self._closed 

136 

137 @property 

138 def close_code(self) -> Optional[int]: 

139 return self._close_code 

140 

141 @property 

142 def protocol(self) -> Optional[str]: 

143 return self._protocol 

144 

145 @property 

146 def compress(self) -> int: 

147 return self._compress 

148 

149 @property 

150 def client_notakeover(self) -> bool: 

151 return self._client_notakeover 

152 

153 def get_extra_info(self, name: str, default: Any = None) -> Any: 

154 """extra info from connection transport""" 

155 conn = self._response.connection 

156 if conn is None: 

157 return default 

158 transport = conn.transport 

159 if transport is None: 

160 return default 

161 return transport.get_extra_info(name, default) 

162 

163 def exception(self) -> Optional[BaseException]: 

164 return self._exception 

165 

166 async def ping(self, message: bytes = b"") -> None: 

167 await self._writer.ping(message) 

168 

169 async def pong(self, message: bytes = b"") -> None: 

170 await self._writer.pong(message) 

171 

172 async def send_str(self, data: str, compress: Optional[int] = None) -> None: 

173 if not isinstance(data, str): 

174 raise TypeError("data argument must be str (%r)" % type(data)) 

175 await self._writer.send(data, binary=False, compress=compress) 

176 

177 async def send_bytes(self, data: bytes, compress: Optional[int] = None) -> None: 

178 if not isinstance(data, (bytes, bytearray, memoryview)): 

179 raise TypeError("data argument must be byte-ish (%r)" % type(data)) 

180 await self._writer.send(data, binary=True, compress=compress) 

181 

182 async def send_json( 

183 self, 

184 data: Any, 

185 compress: Optional[int] = None, 

186 *, 

187 dumps: JSONEncoder = DEFAULT_JSON_ENCODER, 

188 ) -> None: 

189 await self.send_str(dumps(data), compress=compress) 

190 

191 async def close(self, *, code: int = WSCloseCode.OK, message: bytes = b"") -> bool: 

192 # we need to break `receive()` cycle first, 

193 # `close()` may be called from different task 

194 if self._waiting is not None and not self._closed: 

195 self._reader.feed_data(WS_CLOSING_MESSAGE, 0) 

196 await self._waiting 

197 

198 if not self._closed: 

199 self._cancel_heartbeat() 

200 self._closed = True 

201 try: 

202 await self._writer.close(code, message) 

203 except asyncio.CancelledError: 

204 self._close_code = WSCloseCode.ABNORMAL_CLOSURE 

205 self._response.close() 

206 raise 

207 except Exception as exc: 

208 self._close_code = WSCloseCode.ABNORMAL_CLOSURE 

209 self._exception = exc 

210 self._response.close() 

211 return True 

212 

213 if self._closing: 

214 self._response.close() 

215 return True 

216 

217 while True: 

218 try: 

219 async with async_timeout.timeout(self._timeout.ws_close): 

220 msg = await self._reader.read() 

221 except asyncio.CancelledError: 

222 self._close_code = WSCloseCode.ABNORMAL_CLOSURE 

223 self._response.close() 

224 raise 

225 except Exception as exc: 

226 self._close_code = WSCloseCode.ABNORMAL_CLOSURE 

227 self._exception = exc 

228 self._response.close() 

229 return True 

230 

231 if msg.type == WSMsgType.CLOSE: 

232 self._close_code = msg.data 

233 self._response.close() 

234 return True 

235 else: 

236 return False 

237 

238 async def receive(self, timeout: Optional[float] = None) -> WSMessage: 

239 while True: 

240 if self._waiting is not None: 

241 raise RuntimeError("Concurrent call to receive() is not allowed") 

242 

243 if self._closed: 

244 return WS_CLOSED_MESSAGE 

245 elif self._closing: 

246 await self.close() 

247 return WS_CLOSED_MESSAGE 

248 

249 try: 

250 self._waiting = self._loop.create_future() 

251 try: 

252 async with async_timeout.timeout( 

253 timeout or self._timeout.ws_receive 

254 ): 

255 msg = await self._reader.read() 

256 self._reset_heartbeat() 

257 finally: 

258 waiter = self._waiting 

259 self._waiting = None 

260 set_result(waiter, True) 

261 except (asyncio.CancelledError, asyncio.TimeoutError): 

262 self._close_code = WSCloseCode.ABNORMAL_CLOSURE 

263 raise 

264 except EofStream: 

265 self._close_code = WSCloseCode.OK 

266 await self.close() 

267 return WSMessage(WSMsgType.CLOSED, None, None) 

268 except ClientError: 

269 self._closed = True 

270 self._close_code = WSCloseCode.ABNORMAL_CLOSURE 

271 return WS_CLOSED_MESSAGE 

272 except WebSocketError as exc: 

273 self._close_code = exc.code 

274 await self.close(code=exc.code) 

275 return WSMessage(WSMsgType.ERROR, exc, None) 

276 except Exception as exc: 

277 self._exception = exc 

278 self._closing = True 

279 self._close_code = WSCloseCode.ABNORMAL_CLOSURE 

280 await self.close() 

281 return WSMessage(WSMsgType.ERROR, exc, None) 

282 

283 if msg.type == WSMsgType.CLOSE: 

284 self._closing = True 

285 self._close_code = msg.data 

286 # Could be closed elsewhere while awaiting reader 

287 if not self._closed and self._autoclose: # type: ignore[redundant-expr] 

288 await self.close() 

289 elif msg.type == WSMsgType.CLOSING: 

290 self._closing = True 

291 elif msg.type == WSMsgType.PING and self._autoping: 

292 await self.pong(msg.data) 

293 continue 

294 elif msg.type == WSMsgType.PONG and self._autoping: 

295 continue 

296 

297 return msg 

298 

299 async def receive_str(self, *, timeout: Optional[float] = None) -> str: 

300 msg = await self.receive(timeout) 

301 if msg.type != WSMsgType.TEXT: 

302 raise TypeError(f"Received message {msg.type}:{msg.data!r} is not str") 

303 return cast(str, msg.data) 

304 

305 async def receive_bytes(self, *, timeout: Optional[float] = None) -> bytes: 

306 msg = await self.receive(timeout) 

307 if msg.type != WSMsgType.BINARY: 

308 raise TypeError(f"Received message {msg.type}:{msg.data!r} is not bytes") 

309 return cast(bytes, msg.data) 

310 

311 async def receive_json( 

312 self, 

313 *, 

314 loads: JSONDecoder = DEFAULT_JSON_DECODER, 

315 timeout: Optional[float] = None, 

316 ) -> Any: 

317 data = await self.receive_str(timeout=timeout) 

318 return loads(data) 

319 

320 def __aiter__(self) -> "ClientWebSocketResponse": 

321 return self 

322 

323 async def __anext__(self) -> WSMessage: 

324 msg = await self.receive() 

325 if msg.type in (WSMsgType.CLOSE, WSMsgType.CLOSING, WSMsgType.CLOSED): 

326 raise StopAsyncIteration 

327 return msg