Coverage for /pythoncovmergedfiles/medio/medio/src/aiohttp/aiohttp/client_ws.py: 27%
Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1"""WebSocket client for asyncio."""
3import asyncio
4import sys
5from asyncio.base_events import BaseEventLoop
6from collections.abc import Callable
7from types import TracebackType
8from typing import Any, Final, Generic, Literal, overload
10from ._websocket.reader import WebSocketDataQueue, WebSocketReader
11from .client_exceptions import ClientError, ServerTimeoutError, WSMessageTypeError
12from .client_reqrep import ClientResponse
13from .helpers import calculate_timeout_when, frozen_dataclass_decorator, set_result
14from .http import (
15 WS_CLOSED_MESSAGE,
16 WS_CLOSING_MESSAGE,
17 WebSocketError,
18 WSCloseCode,
19 WSMessageDecodeText,
20 WSMessageNoDecodeText,
21 WSMsgType,
22)
23from .http_websocket import _INTERNAL_RECEIVE_TYPES, WebSocketWriter, WSMessageError
24from .streams import EofStream
25from .typedefs import (
26 DEFAULT_JSON_DECODER,
27 DEFAULT_JSON_ENCODER,
28 JSONBytesEncoder,
29 JSONDecoder,
30 JSONEncoder,
31)
33if sys.version_info >= (3, 13):
34 from typing import TypeVar
35else:
36 from typing_extensions import TypeVar
38if sys.version_info >= (3, 11):
39 import asyncio as async_timeout
40 from typing import Self
41else:
42 import async_timeout
43 from typing_extensions import Self
45# TypeVar for whether text messages are decoded to str (True) or kept as bytes (False)
46# Covariant because it only affects return types, not input types
47_DecodeText = TypeVar("_DecodeText", bound=bool, covariant=True, default=Literal[True])
50@frozen_dataclass_decorator
51class ClientWSTimeout:
52 ws_receive: float | None = None
53 ws_close: float | None = None
56DEFAULT_WS_CLIENT_TIMEOUT: Final[ClientWSTimeout] = ClientWSTimeout(
57 ws_receive=None, ws_close=10.0
58)
61class ClientWebSocketResponse(Generic[_DecodeText]):
62 def __init__(
63 self,
64 reader: WebSocketDataQueue,
65 writer: WebSocketWriter,
66 protocol: str | None,
67 response: ClientResponse,
68 timeout: ClientWSTimeout,
69 autoclose: bool,
70 autoping: bool,
71 loop: asyncio.AbstractEventLoop,
72 *,
73 heartbeat: float | None = None,
74 compress: int = 0,
75 client_notakeover: bool = False,
76 ) -> None:
77 self._response = response
78 self._conn = response.connection
80 self._writer = writer
81 self._reader = reader
82 # Set by ClientSession._ws_connect; owns the parser so a stalled
83 # reader parked on the queue by weakref stays alive while this
84 # response can still be drained.
85 self._parser: WebSocketReader | None = None
86 self._protocol = protocol
87 self._closed = False
88 self._closing = False
89 self._close_code: int | None = None
90 self._timeout = timeout
91 self._autoclose = autoclose
92 self._autoping = autoping
93 self._heartbeat = heartbeat
94 self._heartbeat_cb: asyncio.TimerHandle | None = None
95 self._heartbeat_when: float = 0.0
96 if heartbeat is not None:
97 self._pong_heartbeat = heartbeat / 2.0
98 self._pong_response_cb: asyncio.TimerHandle | None = None
99 self._loop = loop
100 self._waiting: bool = False
101 self._close_wait: asyncio.Future[None] | None = None
102 self._exception: BaseException | None = None
103 self._compress = compress
104 self._client_notakeover = client_notakeover
105 self._ping_task: asyncio.Task[None] | None = None
106 self._need_heartbeat_reset = False
107 self._heartbeat_reset_handle: asyncio.Handle | None = None
109 self._reset_heartbeat()
111 def _cancel_heartbeat(self) -> None:
112 self._cancel_pong_response_cb()
113 if self._heartbeat_reset_handle is not None:
114 self._heartbeat_reset_handle.cancel()
115 self._heartbeat_reset_handle = None
116 self._need_heartbeat_reset = False
117 if self._heartbeat_cb is not None:
118 self._heartbeat_cb.cancel()
119 self._heartbeat_cb = None
120 if self._ping_task is not None:
121 self._ping_task.cancel()
122 self._ping_task = None
124 def _cancel_pong_response_cb(self) -> None:
125 if self._pong_response_cb is not None:
126 self._pong_response_cb.cancel()
127 self._pong_response_cb = None
129 def _on_data_received(self) -> None:
130 if self._heartbeat is None or self._need_heartbeat_reset:
131 return
132 loop = self._loop
133 assert loop is not None
134 # Coalesce multiple chunks received in the same loop tick into a single
135 # heartbeat reset. Resetting immediately per chunk increases timer churn.
136 self._need_heartbeat_reset = True
137 self._heartbeat_reset_handle = loop.call_soon(self._flush_heartbeat_reset)
139 def _flush_heartbeat_reset(self) -> None:
140 self._heartbeat_reset_handle = None
141 if not self._need_heartbeat_reset:
142 return
143 self._reset_heartbeat()
144 self._need_heartbeat_reset = False
146 def _reset_heartbeat(self) -> None:
147 if self._heartbeat is None:
148 return
149 self._cancel_pong_response_cb()
150 loop = self._loop
151 assert loop is not None
152 conn = self._conn
153 timeout_ceil_threshold = (
154 conn._connector._timeout_ceil_threshold if conn is not None else 5
155 )
156 now = loop.time()
157 when = calculate_timeout_when(now, self._heartbeat, timeout_ceil_threshold)
158 self._heartbeat_when = when
159 if self._heartbeat_cb is None:
160 # We do not cancel the previous heartbeat_cb here because
161 # it generates a significant amount of TimerHandle churn
162 # which causes asyncio to rebuild the heap frequently.
163 # Instead _send_heartbeat() will reschedule the next
164 # heartbeat if it fires too early.
165 self._heartbeat_cb = loop.call_at(when, self._send_heartbeat)
167 def _send_heartbeat(self) -> None:
168 self._heartbeat_cb = None
170 # If heartbeat reset is pending (data is being received), skip sending
171 # the ping and let the reset callback handle rescheduling the heartbeat.
172 if self._need_heartbeat_reset:
173 return
175 loop = self._loop
176 now = loop.time()
177 if now < self._heartbeat_when:
178 # Heartbeat fired too early, reschedule
179 self._heartbeat_cb = loop.call_at(
180 self._heartbeat_when, self._send_heartbeat
181 )
182 return
184 conn = self._conn
185 timeout_ceil_threshold = (
186 conn._connector._timeout_ceil_threshold if conn is not None else 5
187 )
188 when = calculate_timeout_when(now, self._pong_heartbeat, timeout_ceil_threshold)
189 self._cancel_pong_response_cb()
190 self._pong_response_cb = loop.call_at(when, self._pong_not_received)
192 coro = self._writer.send_frame(b"", WSMsgType.PING)
193 if sys.version_info >= (3, 14):
194 # Try to send the ping immediately to avoid having to schedule
195 # the task on the event loop.
196 if isinstance(loop, BaseEventLoop):
197 ping_task = asyncio.create_task(coro, eager_start=True)
198 else:
199 ping_task = asyncio.Task(coro, loop=loop, eager_start=True)
200 elif sys.version_info >= (3, 12):
201 ping_task = asyncio.Task(coro, loop=loop, eager_start=True)
202 else:
203 ping_task = asyncio.create_task(coro)
205 if not ping_task.done():
206 self._ping_task = ping_task
207 ping_task.add_done_callback(self._ping_task_done)
208 else:
209 self._ping_task_done(ping_task)
211 def _ping_task_done(self, task: "asyncio.Task[None]") -> None:
212 """Callback for when the ping task completes."""
213 if not task.cancelled() and (exc := task.exception()):
214 self._handle_ping_pong_exception(exc)
215 self._ping_task = None
217 def _pong_not_received(self) -> None:
218 self._handle_ping_pong_exception(
219 ServerTimeoutError(f"No PONG received after {self._pong_heartbeat} seconds")
220 )
222 def _handle_ping_pong_exception(self, exc: BaseException) -> None:
223 """Handle exceptions raised during ping/pong processing."""
224 if self._closed:
225 return
226 self._set_closed()
227 # close() is never reached after this; release the parser here.
228 self._parser = None
229 self._close_code = WSCloseCode.ABNORMAL_CLOSURE
230 self._exception = exc
231 self._response.close()
232 if self._waiting and not self._closing:
233 self._reader.feed_data(WSMessageError(data=exc, extra=None))
235 def _set_closed(self) -> None:
236 """Set the connection to closed.
238 Cancel any heartbeat timers and set the closed flag.
239 """
240 self._closed = True
241 self._cancel_heartbeat()
243 def _set_closing(self) -> None:
244 """Set the connection to closing.
246 Cancel any heartbeat timers and set the closing flag.
247 """
248 self._closing = True
249 self._cancel_heartbeat()
251 @property
252 def closed(self) -> bool:
253 return self._closed
255 @property
256 def close_code(self) -> int | None:
257 return self._close_code
259 @property
260 def protocol(self) -> str | None:
261 return self._protocol
263 @property
264 def compress(self) -> int:
265 return self._compress
267 @property
268 def client_notakeover(self) -> bool:
269 return self._client_notakeover
271 def get_extra_info(self, name: str, default: Any = None) -> Any:
272 """extra info from connection transport"""
273 conn = self._response.connection
274 if conn is None:
275 return default
276 transport = conn.transport
277 if transport is None:
278 return default
279 return transport.get_extra_info(name, default)
281 def exception(self) -> BaseException | None:
282 return self._exception
284 async def ping(self, message: bytes = b"") -> None:
285 await self._writer.send_frame(message, WSMsgType.PING)
287 async def pong(self, message: bytes = b"") -> None:
288 await self._writer.send_frame(message, WSMsgType.PONG)
290 async def send_frame(
291 self, message: bytes, opcode: WSMsgType, compress: int | None = None
292 ) -> None:
293 """Send a frame over the websocket."""
294 await self._writer.send_frame(message, opcode, compress)
296 async def send_str(self, data: str, compress: int | None = None) -> None:
297 if not isinstance(data, str):
298 raise TypeError("data argument must be str (%r)" % type(data))
299 await self._writer.send_frame(
300 data.encode("utf-8"), WSMsgType.TEXT, compress=compress
301 )
303 async def send_bytes(self, data: bytes, compress: int | None = None) -> None:
304 if not isinstance(data, (bytes, bytearray, memoryview)):
305 raise TypeError("data argument must be byte-ish (%r)" % type(data))
306 await self._writer.send_frame(data, WSMsgType.BINARY, compress=compress)
308 async def send_json(
309 self,
310 data: Any,
311 compress: int | None = None,
312 *,
313 dumps: JSONEncoder = DEFAULT_JSON_ENCODER,
314 ) -> None:
315 await self.send_str(dumps(data), compress=compress)
317 async def send_json_bytes(
318 self,
319 data: Any,
320 compress: int | None = None,
321 *,
322 dumps: JSONBytesEncoder,
323 ) -> None:
324 """Send JSON data using a bytes-returning encoder as a binary frame.
326 Use this when your JSON encoder (like orjson) returns bytes
327 instead of str, avoiding the encode/decode overhead.
328 """
329 await self.send_bytes(dumps(data), compress=compress)
331 async def close(self, *, code: int = WSCloseCode.OK, message: bytes = b"") -> bool:
332 # we need to break `receive()` cycle first,
333 # `close()` may be called from different task
334 if self._waiting and not self._closing:
335 assert self._loop is not None
336 self._close_wait = self._loop.create_future()
337 self._set_closing()
338 self._reader.feed_data(WS_CLOSING_MESSAGE)
339 await self._close_wait
341 if self._closed:
342 return False
344 self._set_closed()
345 try:
346 try:
347 await self._writer.close(code, message)
348 except asyncio.CancelledError:
349 self._close_code = WSCloseCode.ABNORMAL_CLOSURE
350 self._response.close()
351 raise
352 except Exception as exc:
353 self._close_code = WSCloseCode.ABNORMAL_CLOSURE
354 self._exception = exc
355 self._response.close()
356 return True
358 if self._close_code:
359 self._response.close()
360 return True
362 while True:
363 try:
364 async with async_timeout.timeout(self._timeout.ws_close):
365 msg = await self._reader.read()
366 except asyncio.CancelledError:
367 self._close_code = WSCloseCode.ABNORMAL_CLOSURE
368 self._response.close()
369 raise
370 except Exception as exc:
371 self._close_code = WSCloseCode.ABNORMAL_CLOSURE
372 self._exception = exc
373 self._response.close()
374 return True
376 if msg.type is WSMsgType.CLOSE:
377 self._close_code = msg.data
378 self._response.close()
379 return True
380 finally:
381 # Once closed the response can no longer be drained; release the
382 # parser and the stash it retains.
383 self._parser = None
385 @overload
386 async def receive(
387 self: "ClientWebSocketResponse[Literal[True]]", timeout: float | None = None
388 ) -> WSMessageDecodeText: ...
390 @overload
391 async def receive(
392 self: "ClientWebSocketResponse[Literal[False]]", timeout: float | None = None
393 ) -> WSMessageNoDecodeText: ...
395 @overload
396 async def receive(
397 self: "ClientWebSocketResponse[_DecodeText]", timeout: float | None = None
398 ) -> WSMessageDecodeText | WSMessageNoDecodeText: ...
400 async def receive(
401 self, timeout: float | None = None
402 ) -> WSMessageDecodeText | WSMessageNoDecodeText:
403 receive_timeout = timeout or self._timeout.ws_receive
405 while True:
406 if self._waiting:
407 raise RuntimeError("Concurrent call to receive() is not allowed")
409 if self._closed:
410 return WS_CLOSED_MESSAGE
411 elif self._closing:
412 await self.close()
413 return WS_CLOSED_MESSAGE
415 try:
416 self._waiting = True
417 try:
418 if receive_timeout:
419 # Entering the context manager and creating
420 # Timeout() object can take almost 50% of the
421 # run time in this loop so we avoid it if
422 # there is no read timeout.
423 async with async_timeout.timeout(receive_timeout):
424 msg = await self._reader.read()
425 else:
426 msg = await self._reader.read()
427 finally:
428 self._waiting = False
429 if self._close_wait:
430 set_result(self._close_wait, None)
431 except (asyncio.CancelledError, asyncio.TimeoutError):
432 self._close_code = WSCloseCode.ABNORMAL_CLOSURE
433 raise
434 except EofStream:
435 self._close_code = WSCloseCode.OK
436 await self.close()
437 return WS_CLOSED_MESSAGE
438 except ClientError:
439 # Likely ServerDisconnectedError when connection is lost.
440 # close() is not called on this path, so release the parser
441 # and the stash it retains here.
442 self._parser = None
443 self._set_closed()
444 self._close_code = WSCloseCode.ABNORMAL_CLOSURE
445 return WS_CLOSED_MESSAGE
446 except WebSocketError as exc:
447 self._close_code = exc.code
448 await self.close(code=exc.code)
449 return WSMessageError(data=exc)
450 except Exception as exc:
451 self._exception = exc
452 self._set_closing()
453 self._close_code = WSCloseCode.ABNORMAL_CLOSURE
454 await self.close()
455 return WSMessageError(data=exc)
457 if msg.type not in _INTERNAL_RECEIVE_TYPES:
458 # If its not a close/closing/ping/pong message
459 # we can return it immediately
460 return msg
462 if msg.type is WSMsgType.CLOSE:
463 self._set_closing()
464 self._close_code = msg.data
465 # Could be closed elsewhere while awaiting reader
466 if not self._closed and self._autoclose: # type: ignore[redundant-expr]
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
476 return msg
478 @overload
479 async def receive_str(
480 self: "ClientWebSocketResponse[Literal[True]]", *, timeout: float | None = None
481 ) -> str: ...
483 @overload
484 async def receive_str(
485 self: "ClientWebSocketResponse[Literal[False]]", *, timeout: float | None = None
486 ) -> bytes: ...
488 @overload
489 async def receive_str(
490 self: "ClientWebSocketResponse[_DecodeText]", *, timeout: float | None = None
491 ) -> str | bytes: ...
493 async def receive_str(self, *, timeout: float | None = None) -> str | bytes:
494 """Receive TEXT message.
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 msg.data
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 msg.data
513 @overload
514 async def receive_json(
515 self: "ClientWebSocketResponse[Literal[True]]",
516 *,
517 loads: JSONDecoder = ...,
518 timeout: float | None = None,
519 ) -> Any: ...
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: ...
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: ...
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]
546 def __aiter__(self) -> Self:
547 return self
549 @overload
550 async def __anext__(
551 self: "ClientWebSocketResponse[Literal[True]]",
552 ) -> WSMessageDecodeText: ...
554 @overload
555 async def __anext__(
556 self: "ClientWebSocketResponse[Literal[False]]",
557 ) -> WSMessageNoDecodeText: ...
559 @overload
560 async def __anext__(
561 self: "ClientWebSocketResponse[_DecodeText]",
562 ) -> WSMessageDecodeText | WSMessageNoDecodeText: ...
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
570 async def __aenter__(self) -> Self:
571 return self
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()