Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/aiohttp/client_ws.py: 27%

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

2 

3import asyncio 

4import sys 

5from asyncio.base_events import BaseEventLoop 

6from collections.abc import Callable 

7from types import TracebackType 

8from typing import Any, Generic, Literal, Optional, cast, overload 

9 

10import attr 

11 

12from ._websocket.reader import WebSocketDataQueue, WebSocketReader 

13from .client_exceptions import ClientError, ServerTimeoutError, WSMessageTypeError 

14from .client_reqrep import ClientResponse 

15from .helpers import calculate_timeout_when, set_result 

16from .http import ( 

17 WS_CLOSED_MESSAGE, 

18 WS_CLOSING_MESSAGE, 

19 WebSocketError, 

20 WSCloseCode, 

21 WSMessage, 

22 WSMessageDecodeText, 

23 WSMessageNoDecodeText, 

24 WSMsgType, 

25) 

26from .http_websocket import _INTERNAL_RECEIVE_TYPES, WebSocketWriter 

27from .streams import EofStream 

28from .typedefs import ( 

29 DEFAULT_JSON_DECODER, 

30 DEFAULT_JSON_ENCODER, 

31 JSONBytesEncoder, 

32 JSONDecoder, 

33 JSONEncoder, 

34) 

35 

36if sys.version_info >= (3, 13): 

37 from typing import TypeVar 

38else: 

39 from typing_extensions import TypeVar 

40 

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

42 import asyncio as async_timeout 

43 from typing import Self 

44else: 

45 import async_timeout 

46 from typing_extensions import Self 

47 

48# TypeVar for whether text messages are decoded to str (True) or kept as bytes (False) 

49# Covariant because it only affects return types, not input types 

50_DecodeText = TypeVar("_DecodeText", bound=bool, covariant=True, default=Literal[True]) 

51 

52 

53@attr.s(frozen=True, slots=True) 

54class ClientWSTimeout: 

55 ws_receive = attr.ib(type=Optional[float], default=None) 

56 ws_close = attr.ib(type=Optional[float], default=None) 

57 

58 

59DEFAULT_WS_CLIENT_TIMEOUT = ClientWSTimeout(ws_receive=None, ws_close=10.0) 

60 

61 

62class ClientWebSocketResponse(Generic[_DecodeText]): 

63 def __init__( 

64 self, 

65 reader: WebSocketDataQueue, 

66 writer: WebSocketWriter, 

67 protocol: str | None, 

68 response: ClientResponse, 

69 timeout: ClientWSTimeout, 

70 autoclose: bool, 

71 autoping: bool, 

72 loop: asyncio.AbstractEventLoop, 

73 *, 

74 heartbeat: float | None = None, 

75 compress: int = 0, 

76 client_notakeover: bool = False, 

77 ) -> None: 

78 self._response = response 

79 self._conn = response.connection 

80 

81 self._writer = writer 

82 self._reader = reader 

83 # Set by ClientSession._ws_connect; owns the parser so a stalled 

84 # reader parked on the queue by weakref stays alive while this 

85 # response can still be drained. 

86 self._parser: WebSocketReader | None = None 

87 self._protocol = protocol 

88 self._closed = False 

89 self._closing = False 

90 self._close_code: int | None = None 

91 self._timeout = timeout 

92 self._autoclose = autoclose 

93 self._autoping = autoping 

94 self._heartbeat = heartbeat 

95 self._heartbeat_cb: asyncio.TimerHandle | None = None 

96 self._heartbeat_when: float = 0.0 

97 if heartbeat is not None: 

98 self._pong_heartbeat = heartbeat / 2.0 

99 self._pong_response_cb: asyncio.TimerHandle | None = None 

100 self._loop = loop 

101 self._waiting: bool = False 

102 self._close_wait: asyncio.Future[None] | None = None 

103 self._exception: BaseException | None = None 

104 self._compress = compress 

105 self._client_notakeover = client_notakeover 

106 self._ping_task: asyncio.Task[None] | None = None 

107 self._need_heartbeat_reset = False 

108 self._heartbeat_reset_handle: asyncio.Handle | None = None 

109 

110 self._reset_heartbeat() 

111 

112 def _cancel_heartbeat(self) -> None: 

113 self._cancel_pong_response_cb() 

114 if self._heartbeat_reset_handle is not None: 

115 self._heartbeat_reset_handle.cancel() 

116 self._heartbeat_reset_handle = None 

117 self._need_heartbeat_reset = False 

118 if self._heartbeat_cb is not None: 

119 self._heartbeat_cb.cancel() 

120 self._heartbeat_cb = None 

121 if self._ping_task is not None: 

122 self._ping_task.cancel() 

123 self._ping_task = None 

124 

125 def _cancel_pong_response_cb(self) -> None: 

126 if self._pong_response_cb is not None: 

127 self._pong_response_cb.cancel() 

128 self._pong_response_cb = None 

129 

130 def _on_data_received(self) -> None: 

131 if self._heartbeat is None or self._need_heartbeat_reset: 

132 return 

133 loop = self._loop 

134 assert loop is not None 

135 # Coalesce multiple chunks received in the same loop tick into a single 

136 # heartbeat reset. Resetting immediately per chunk increases timer churn. 

137 self._need_heartbeat_reset = True 

138 self._heartbeat_reset_handle = loop.call_soon(self._flush_heartbeat_reset) 

139 

140 def _flush_heartbeat_reset(self) -> None: 

141 self._heartbeat_reset_handle = None 

142 if not self._need_heartbeat_reset: 

143 return 

144 self._reset_heartbeat() 

145 self._need_heartbeat_reset = False 

146 

147 def _reset_heartbeat(self) -> None: 

148 if self._heartbeat is None: 

149 return 

150 self._cancel_pong_response_cb() 

151 loop = self._loop 

152 assert loop is not None 

153 conn = self._conn 

154 timeout_ceil_threshold = ( 

155 conn._connector._timeout_ceil_threshold if conn is not None else 5 

156 ) 

157 now = loop.time() 

158 when = calculate_timeout_when(now, self._heartbeat, timeout_ceil_threshold) 

159 self._heartbeat_when = when 

160 if self._heartbeat_cb is None: 

161 # We do not cancel the previous heartbeat_cb here because 

162 # it generates a significant amount of TimerHandle churn 

163 # which causes asyncio to rebuild the heap frequently. 

164 # Instead _send_heartbeat() will reschedule the next 

165 # heartbeat if it fires too early. 

166 self._heartbeat_cb = loop.call_at(when, self._send_heartbeat) 

167 

168 def _send_heartbeat(self) -> None: 

169 self._heartbeat_cb = None 

170 

171 # If heartbeat reset is pending (data is being received), skip sending 

172 # the ping and let the reset callback handle rescheduling the heartbeat. 

173 if self._need_heartbeat_reset: 

174 return 

175 

176 loop = self._loop 

177 now = loop.time() 

178 if now < self._heartbeat_when: 

179 # Heartbeat fired too early, reschedule 

180 self._heartbeat_cb = loop.call_at( 

181 self._heartbeat_when, self._send_heartbeat 

182 ) 

183 return 

184 

185 conn = self._conn 

186 timeout_ceil_threshold = ( 

187 conn._connector._timeout_ceil_threshold if conn is not None else 5 

188 ) 

189 when = calculate_timeout_when(now, self._pong_heartbeat, timeout_ceil_threshold) 

190 self._cancel_pong_response_cb() 

191 self._pong_response_cb = loop.call_at(when, self._pong_not_received) 

192 

193 coro = self._writer.send_frame(b"", WSMsgType.PING) 

194 if sys.version_info >= (3, 14): 

195 # Try to send the ping immediately to avoid having to schedule 

196 # the task on the event loop. 

197 if isinstance(loop, BaseEventLoop): 

198 ping_task = asyncio.create_task(coro, eager_start=True) 

199 else: 

200 ping_task = asyncio.Task(coro, loop=loop, eager_start=True) 

201 elif sys.version_info >= (3, 12): 

202 ping_task = asyncio.Task(coro, loop=loop, eager_start=True) 

203 else: 

204 ping_task = asyncio.create_task(coro) 

205 

206 if not ping_task.done(): 

207 self._ping_task = ping_task 

208 ping_task.add_done_callback(self._ping_task_done) 

209 else: 

210 self._ping_task_done(ping_task) 

211 

212 def _ping_task_done(self, task: "asyncio.Task[None]") -> None: 

213 """Callback for when the ping task completes.""" 

214 if not task.cancelled() and (exc := task.exception()): 

215 self._handle_ping_pong_exception(exc) 

216 self._ping_task = None 

217 

218 def _pong_not_received(self) -> None: 

219 self._handle_ping_pong_exception( 

220 ServerTimeoutError(f"No PONG received after {self._pong_heartbeat} seconds") 

221 ) 

222 

223 def _handle_ping_pong_exception(self, exc: BaseException) -> None: 

224 """Handle exceptions raised during ping/pong processing.""" 

225 if self._closed: 

226 return 

227 self._set_closed() 

228 # close() is never reached after this; release the parser here. 

229 self._parser = None 

230 self._close_code = WSCloseCode.ABNORMAL_CLOSURE 

231 self._exception = exc 

232 self._response.close() 

233 if self._waiting and not self._closing: 

234 self._reader.feed_data(WSMessage(WSMsgType.ERROR, exc, None), 0) 

235 

236 def _set_closed(self) -> None: 

237 """Set the connection to closed. 

238 

239 Cancel any heartbeat timers and set the closed flag. 

240 """ 

241 self._closed = True 

242 self._cancel_heartbeat() 

243 

244 def _set_closing(self) -> None: 

245 """Set the connection to closing. 

246 

247 Cancel any heartbeat timers and set the closing flag. 

248 """ 

249 self._closing = True 

250 self._cancel_heartbeat() 

251 

252 @property 

253 def closed(self) -> bool: 

254 return self._closed 

255 

256 @property 

257 def close_code(self) -> int | None: 

258 return self._close_code 

259 

260 @property 

261 def protocol(self) -> str | None: 

262 return self._protocol 

263 

264 @property 

265 def compress(self) -> int: 

266 return self._compress 

267 

268 @property 

269 def client_notakeover(self) -> bool: 

270 return self._client_notakeover 

271 

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

273 """extra info from connection transport""" 

274 conn = self._response.connection 

275 if conn is None: 

276 return default 

277 transport = conn.transport 

278 if transport is None: 

279 return default 

280 return transport.get_extra_info(name, default) 

281 

282 def exception(self) -> BaseException | None: 

283 return self._exception 

284 

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

286 await self._writer.send_frame(message, WSMsgType.PING) 

287 

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

289 await self._writer.send_frame(message, WSMsgType.PONG) 

290 

291 async def send_frame( 

292 self, message: bytes, opcode: WSMsgType, compress: int | None = None 

293 ) -> None: 

294 """Send a frame over the websocket.""" 

295 await self._writer.send_frame(message, opcode, compress) 

296 

297 async def send_str(self, data: str, compress: int | None = None) -> None: 

298 if not isinstance(data, str): 

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

300 await self._writer.send_frame( 

301 data.encode("utf-8"), WSMsgType.TEXT, compress=compress 

302 ) 

303 

304 async def send_bytes(self, data: bytes, compress: int | None = None) -> None: 

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

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

307 await self._writer.send_frame(data, WSMsgType.BINARY, compress=compress) 

308 

309 async def send_json( 

310 self, 

311 data: Any, 

312 compress: int | None = None, 

313 *, 

314 dumps: JSONEncoder = DEFAULT_JSON_ENCODER, 

315 ) -> None: 

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

317 

318 async def send_json_bytes( 

319 self, 

320 data: Any, 

321 compress: int | None = None, 

322 *, 

323 dumps: JSONBytesEncoder, 

324 ) -> None: 

325 """Send JSON data using a bytes-returning encoder as a binary frame. 

326 

327 Use this when your JSON encoder (like orjson) returns bytes 

328 instead of str, avoiding the encode/decode overhead. 

329 """ 

330 await self.send_bytes(dumps(data), compress=compress) 

331 

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

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

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

335 if self._waiting and not self._closing: 

336 assert self._loop is not None 

337 self._close_wait = self._loop.create_future() 

338 self._set_closing() 

339 self._reader.feed_data(WS_CLOSING_MESSAGE, 0) 

340 await self._close_wait 

341 

342 if self._closed: 

343 return False 

344 

345 self._set_closed() 

346 try: 

347 try: 

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

349 except asyncio.CancelledError: 

350 self._close_code = WSCloseCode.ABNORMAL_CLOSURE 

351 self._response.close() 

352 raise 

353 except Exception as exc: 

354 self._close_code = WSCloseCode.ABNORMAL_CLOSURE 

355 self._exception = exc 

356 self._response.close() 

357 return True 

358 

359 if self._close_code: 

360 self._response.close() 

361 return True 

362 

363 while True: 

364 try: 

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

366 msg = await self._reader.read() 

367 except asyncio.CancelledError: 

368 self._close_code = WSCloseCode.ABNORMAL_CLOSURE 

369 self._response.close() 

370 raise 

371 except Exception as exc: 

372 self._close_code = WSCloseCode.ABNORMAL_CLOSURE 

373 self._exception = exc 

374 self._response.close() 

375 return True 

376 

377 if msg.type is WSMsgType.CLOSE: 

378 self._close_code = msg.data 

379 self._response.close() 

380 return True 

381 finally: 

382 # Once closed the response can no longer be drained; release the 

383 # parser and the stash it retains. 

384 self._parser = None 

385 

386 @overload 

387 async def receive( 

388 self: "ClientWebSocketResponse[Literal[True]]", timeout: float | None = None 

389 ) -> WSMessageDecodeText: ... 

390 

391 @overload 

392 async def receive( 

393 self: "ClientWebSocketResponse[Literal[False]]", timeout: float | None = None 

394 ) -> WSMessageNoDecodeText: ... 

395 

396 @overload 

397 async def receive( 

398 self: "ClientWebSocketResponse[_DecodeText]", timeout: float | None = None 

399 ) -> WSMessageDecodeText | WSMessageNoDecodeText: ... 

400 

401 async def receive( 

402 self, timeout: float | None = None 

403 ) -> WSMessageDecodeText | WSMessageNoDecodeText: 

404 receive_timeout = timeout or self._timeout.ws_receive 

405 

406 while True: 

407 if self._waiting: 

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

409 

410 if self._closed: 

411 return WS_CLOSED_MESSAGE 

412 elif self._closing: 

413 await self.close() 

414 return WS_CLOSED_MESSAGE 

415 

416 try: 

417 self._waiting = True 

418 try: 

419 if receive_timeout: 

420 # Entering the context manager and creating 

421 # Timeout() object can take almost 50% of the 

422 # run time in this loop so we avoid it if 

423 # there is no read timeout. 

424 async with async_timeout.timeout(receive_timeout): 

425 msg = await self._reader.read() 

426 else: 

427 msg = await self._reader.read() 

428 finally: 

429 self._waiting = False 

430 if self._close_wait: 

431 set_result(self._close_wait, None) 

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

433 self._close_code = WSCloseCode.ABNORMAL_CLOSURE 

434 raise 

435 except EofStream: 

436 self._close_code = WSCloseCode.OK 

437 await self.close() 

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

439 except ClientError: 

440 # Likely ServerDisconnectedError when connection is lost. 

441 # close() is not called on this path, so release the parser 

442 # and the stash it retains here. 

443 self._parser = None 

444 self._set_closed() 

445 self._close_code = WSCloseCode.ABNORMAL_CLOSURE 

446 return WS_CLOSED_MESSAGE 

447 except WebSocketError as exc: 

448 self._close_code = exc.code 

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

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

451 except Exception as exc: 

452 self._exception = exc 

453 self._set_closing() 

454 self._close_code = WSCloseCode.ABNORMAL_CLOSURE 

455 await self.close() 

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

457 

458 if msg.type not in _INTERNAL_RECEIVE_TYPES: 

459 # If its not a close/closing/ping/pong message 

460 # we can return it immediately 

461 return msg 

462 

463 if msg.type is WSMsgType.CLOSE: 

464 self._set_closing() 

465 self._close_code = msg.data 

466 if not self._closed and self._autoclose: 

467 await self.close() 

468 elif msg.type is WSMsgType.CLOSING: 

469 self._set_closing() 

470 elif msg.type is WSMsgType.PING and self._autoping: 

471 await self.pong(msg.data) 

472 continue 

473 elif msg.type is WSMsgType.PONG and self._autoping: 

474 continue 

475 

476 return msg 

477 

478 @overload 

479 async def receive_str( 

480 self: "ClientWebSocketResponse[Literal[True]]", *, timeout: float | None = None 

481 ) -> str: ... 

482 

483 @overload 

484 async def receive_str( 

485 self: "ClientWebSocketResponse[Literal[False]]", *, timeout: float | None = None 

486 ) -> bytes: ... 

487 

488 @overload 

489 async def receive_str( 

490 self: "ClientWebSocketResponse[_DecodeText]", *, timeout: float | None = None 

491 ) -> str | bytes: ... 

492 

493 async def receive_str(self, *, timeout: float | None = None) -> str | bytes: 

494 """Receive TEXT message. 

495 

496 Returns str when decode_text=True (default), bytes when decode_text=False. 

497 """ 

498 msg = await self.receive(timeout) 

499 if msg.type is not WSMsgType.TEXT: 

500 raise WSMessageTypeError( 

501 f"Received message {msg.type}:{msg.data!r} is not WSMsgType.TEXT" 

502 ) 

503 return cast(str, msg.data) 

504 

505 async def receive_bytes(self, *, timeout: float | None = None) -> bytes: 

506 msg = await self.receive(timeout) 

507 if msg.type is not WSMsgType.BINARY: 

508 raise WSMessageTypeError( 

509 f"Received message {msg.type}:{msg.data!r} is not WSMsgType.BINARY" 

510 ) 

511 return cast(bytes, msg.data) 

512 

513 @overload 

514 async def receive_json( 

515 self: "ClientWebSocketResponse[Literal[True]]", 

516 *, 

517 loads: JSONDecoder = ..., 

518 timeout: float | None = None, 

519 ) -> Any: ... 

520 

521 @overload 

522 async def receive_json( 

523 self: "ClientWebSocketResponse[Literal[False]]", 

524 *, 

525 loads: Callable[[bytes], Any] = ..., 

526 timeout: float | None = None, 

527 ) -> Any: ... 

528 

529 @overload 

530 async def receive_json( 

531 self: "ClientWebSocketResponse[_DecodeText]", 

532 *, 

533 loads: JSONDecoder | Callable[[bytes], Any] = ..., 

534 timeout: float | None = None, 

535 ) -> Any: ... 

536 

537 async def receive_json( 

538 self, 

539 *, 

540 loads: JSONDecoder | Callable[[bytes], Any] = DEFAULT_JSON_DECODER, 

541 timeout: float | None = None, 

542 ) -> Any: 

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

544 return loads(data) # type: ignore[arg-type] 

545 

546 def __aiter__(self) -> Self: 

547 return self 

548 

549 @overload 

550 async def __anext__( 

551 self: "ClientWebSocketResponse[Literal[True]]", 

552 ) -> WSMessageDecodeText: ... 

553 

554 @overload 

555 async def __anext__( 

556 self: "ClientWebSocketResponse[Literal[False]]", 

557 ) -> WSMessageNoDecodeText: ... 

558 

559 @overload 

560 async def __anext__( 

561 self: "ClientWebSocketResponse[_DecodeText]", 

562 ) -> WSMessageDecodeText | WSMessageNoDecodeText: ... 

563 

564 async def __anext__(self) -> WSMessageDecodeText | WSMessageNoDecodeText: 

565 msg = await self.receive() 

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

567 raise StopAsyncIteration 

568 return msg 

569 

570 async def __aenter__(self) -> Self: 

571 return self 

572 

573 async def __aexit__( 

574 self, 

575 exc_type: type[BaseException] | None, 

576 exc_val: BaseException | None, 

577 exc_tb: TracebackType | None, 

578 ) -> None: 

579 await self.close()