Coverage for /pythoncovmergedfiles/medio/medio/src/aiohttp/aiohttp/client_ws.py: 26%
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, Self, 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
38# TypeVar for whether text messages are decoded to str (True) or kept as bytes (False)
39# Covariant because it only affects return types, not input types
40_DecodeText = TypeVar("_DecodeText", bound=bool, covariant=True, default=Literal[True])
43@frozen_dataclass_decorator
44class ClientWSTimeout:
45 ws_receive: float | None = None
46 ws_close: float | None = None
49DEFAULT_WS_CLIENT_TIMEOUT: Final[ClientWSTimeout] = ClientWSTimeout(
50 ws_receive=None, ws_close=10.0
51)
54class ClientWebSocketResponse(Generic[_DecodeText]):
55 def __init__(
56 self,
57 reader: WebSocketDataQueue,
58 writer: WebSocketWriter,
59 protocol: str | None,
60 response: ClientResponse,
61 timeout: ClientWSTimeout,
62 autoclose: bool,
63 autoping: bool,
64 loop: asyncio.AbstractEventLoop,
65 *,
66 heartbeat: float | None = None,
67 compress: int = 0,
68 client_notakeover: bool = False,
69 ) -> None:
70 self._response = response
71 self._conn = response.connection
73 self._writer = writer
74 self._reader = reader
75 # Set by ClientSession._ws_connect; owns the parser so a stalled
76 # reader parked on the queue by weakref stays alive while this
77 # response can still be drained.
78 self._parser: WebSocketReader | None = None
79 self._protocol = protocol
80 self._closed = False
81 self._closing = False
82 self._close_code: int | None = None
83 self._timeout = timeout
84 self._autoclose = autoclose
85 self._autoping = autoping
86 self._heartbeat = heartbeat
87 self._heartbeat_cb: asyncio.TimerHandle | None = None
88 self._heartbeat_when: float = 0.0
89 if heartbeat is not None:
90 self._pong_heartbeat = heartbeat / 2.0
91 self._pong_response_cb: asyncio.TimerHandle | None = None
92 self._loop = loop
93 self._waiting: bool = False
94 self._close_wait: asyncio.Future[None] | None = None
95 self._exception: BaseException | None = None
96 self._compress = compress
97 self._client_notakeover = client_notakeover
98 self._ping_task: asyncio.Task[None] | None = None
99 self._need_heartbeat_reset = False
100 self._heartbeat_reset_handle: asyncio.Handle | None = None
102 self._reset_heartbeat()
104 def _cancel_heartbeat(self) -> None:
105 self._cancel_pong_response_cb()
106 if self._heartbeat_reset_handle is not None:
107 self._heartbeat_reset_handle.cancel()
108 self._heartbeat_reset_handle = None
109 self._need_heartbeat_reset = False
110 if self._heartbeat_cb is not None:
111 self._heartbeat_cb.cancel()
112 self._heartbeat_cb = None
113 if self._ping_task is not None:
114 self._ping_task.cancel()
115 self._ping_task = None
117 def _cancel_pong_response_cb(self) -> None:
118 if self._pong_response_cb is not None:
119 self._pong_response_cb.cancel()
120 self._pong_response_cb = None
122 def _on_data_received(self) -> None:
123 if self._heartbeat is None or self._need_heartbeat_reset:
124 return
125 loop = self._loop
126 assert loop is not None
127 # Coalesce multiple chunks received in the same loop tick into a single
128 # heartbeat reset. Resetting immediately per chunk increases timer churn.
129 self._need_heartbeat_reset = True
130 self._heartbeat_reset_handle = loop.call_soon(self._flush_heartbeat_reset)
132 def _flush_heartbeat_reset(self) -> None:
133 self._heartbeat_reset_handle = None
134 if not self._need_heartbeat_reset:
135 return
136 self._reset_heartbeat()
137 self._need_heartbeat_reset = False
139 def _reset_heartbeat(self) -> None:
140 if self._heartbeat is None:
141 return
142 self._cancel_pong_response_cb()
143 loop = self._loop
144 assert loop is not None
145 conn = self._conn
146 timeout_ceil_threshold = (
147 conn._connector._timeout_ceil_threshold if conn is not None else 5
148 )
149 now = loop.time()
150 when = calculate_timeout_when(now, self._heartbeat, timeout_ceil_threshold)
151 self._heartbeat_when = when
152 if self._heartbeat_cb is None:
153 # We do not cancel the previous heartbeat_cb here because
154 # it generates a significant amount of TimerHandle churn
155 # which causes asyncio to rebuild the heap frequently.
156 # Instead _send_heartbeat() will reschedule the next
157 # heartbeat if it fires too early.
158 self._heartbeat_cb = loop.call_at(when, self._send_heartbeat)
160 def _send_heartbeat(self) -> None:
161 self._heartbeat_cb = None
163 # If heartbeat reset is pending (data is being received), skip sending
164 # the ping and let the reset callback handle rescheduling the heartbeat.
165 if self._need_heartbeat_reset:
166 return
168 loop = self._loop
169 now = loop.time()
170 if now < self._heartbeat_when:
171 # Heartbeat fired too early, reschedule
172 self._heartbeat_cb = loop.call_at(
173 self._heartbeat_when, self._send_heartbeat
174 )
175 return
177 conn = self._conn
178 timeout_ceil_threshold = (
179 conn._connector._timeout_ceil_threshold if conn is not None else 5
180 )
181 when = calculate_timeout_when(now, self._pong_heartbeat, timeout_ceil_threshold)
182 self._cancel_pong_response_cb()
183 self._pong_response_cb = loop.call_at(when, self._pong_not_received)
185 coro = self._writer.send_frame(b"", WSMsgType.PING)
186 if sys.version_info >= (3, 14):
187 # Try to send the ping immediately to avoid having to schedule
188 # the task on the event loop.
189 if isinstance(loop, BaseEventLoop):
190 ping_task = asyncio.create_task(coro, eager_start=True)
191 else:
192 ping_task = asyncio.Task(coro, loop=loop, eager_start=True)
193 elif sys.version_info >= (3, 12):
194 ping_task = asyncio.Task(coro, loop=loop, eager_start=True)
195 else:
196 ping_task = asyncio.create_task(coro)
198 if not ping_task.done():
199 self._ping_task = ping_task
200 ping_task.add_done_callback(self._ping_task_done)
201 else:
202 self._ping_task_done(ping_task)
204 def _ping_task_done(self, task: "asyncio.Task[None]") -> None:
205 """Callback for when the ping task completes."""
206 if not task.cancelled() and (exc := task.exception()):
207 self._handle_ping_pong_exception(exc)
208 self._ping_task = None
210 def _pong_not_received(self) -> None:
211 self._handle_ping_pong_exception(
212 ServerTimeoutError(f"No PONG received after {self._pong_heartbeat} seconds")
213 )
215 def _handle_ping_pong_exception(self, exc: BaseException) -> None:
216 """Handle exceptions raised during ping/pong processing."""
217 if self._closed:
218 return
219 self._set_closed()
220 # close() is never reached after this; release the parser here.
221 self._parser = None
222 self._close_code = WSCloseCode.ABNORMAL_CLOSURE
223 self._exception = exc
224 self._response.close()
225 if self._waiting and not self._closing:
226 self._reader.feed_data(WSMessageError(data=exc, extra=None))
228 def _set_closed(self) -> None:
229 """Set the connection to closed.
231 Cancel any heartbeat timers and set the closed flag.
232 """
233 self._closed = True
234 self._cancel_heartbeat()
236 def _set_closing(self) -> None:
237 """Set the connection to closing.
239 Cancel any heartbeat timers and set the closing flag.
240 """
241 self._closing = True
242 self._cancel_heartbeat()
244 @property
245 def closed(self) -> bool:
246 return self._closed
248 @property
249 def close_code(self) -> int | None:
250 return self._close_code
252 @property
253 def protocol(self) -> str | None:
254 return self._protocol
256 @property
257 def compress(self) -> int:
258 return self._compress
260 @property
261 def client_notakeover(self) -> bool:
262 return self._client_notakeover
264 def get_extra_info(self, name: str, default: Any = None) -> Any:
265 """extra info from connection transport"""
266 conn = self._response.connection
267 if conn is None:
268 return default
269 transport = conn.transport
270 if transport is None:
271 return default
272 return transport.get_extra_info(name, default)
274 def exception(self) -> BaseException | None:
275 return self._exception
277 async def ping(self, message: bytes = b"") -> None:
278 await self._writer.send_frame(message, WSMsgType.PING)
280 async def pong(self, message: bytes = b"") -> None:
281 await self._writer.send_frame(message, WSMsgType.PONG)
283 async def send_frame(
284 self, message: bytes, opcode: WSMsgType, compress: int | None = None
285 ) -> None:
286 """Send a frame over the websocket."""
287 await self._writer.send_frame(message, opcode, compress)
289 async def send_str(self, data: str, compress: int | None = None) -> None:
290 if not isinstance(data, str):
291 raise TypeError("data argument must be str (%r)" % type(data))
292 await self._writer.send_frame(
293 data.encode("utf-8"), WSMsgType.TEXT, compress=compress
294 )
296 async def send_bytes(self, data: bytes, compress: int | None = None) -> None:
297 if not isinstance(data, (bytes, bytearray, memoryview)):
298 raise TypeError("data argument must be byte-ish (%r)" % type(data))
299 await self._writer.send_frame(data, WSMsgType.BINARY, compress=compress)
301 async def send_json(
302 self,
303 data: Any,
304 compress: int | None = None,
305 *,
306 dumps: JSONEncoder = DEFAULT_JSON_ENCODER,
307 ) -> None:
308 await self.send_str(dumps(data), compress=compress)
310 async def send_json_bytes(
311 self,
312 data: Any,
313 compress: int | None = None,
314 *,
315 dumps: JSONBytesEncoder,
316 ) -> None:
317 """Send JSON data using a bytes-returning encoder as a binary frame.
319 Use this when your JSON encoder (like orjson) returns bytes
320 instead of str, avoiding the encode/decode overhead.
321 """
322 await self.send_bytes(dumps(data), compress=compress)
324 async def close(self, *, code: int = WSCloseCode.OK, message: bytes = b"") -> bool:
325 # we need to break `receive()` cycle first,
326 # `close()` may be called from different task
327 if self._waiting and not self._closing:
328 assert self._loop is not None
329 self._close_wait = self._loop.create_future()
330 self._set_closing()
331 self._reader.feed_data(WS_CLOSING_MESSAGE)
332 await self._close_wait
334 if self._closed:
335 return False
337 self._set_closed()
338 try:
339 try:
340 await self._writer.close(code, message)
341 except asyncio.CancelledError:
342 self._close_code = WSCloseCode.ABNORMAL_CLOSURE
343 self._response.close()
344 raise
345 except Exception as exc:
346 self._close_code = WSCloseCode.ABNORMAL_CLOSURE
347 self._exception = exc
348 self._response.close()
349 return True
351 if self._close_code:
352 self._response.close()
353 return True
355 while True:
356 try:
357 async with asyncio.timeout(self._timeout.ws_close):
358 msg = await self._reader.read()
359 except asyncio.CancelledError:
360 self._close_code = WSCloseCode.ABNORMAL_CLOSURE
361 self._response.close()
362 raise
363 except Exception as exc:
364 self._close_code = WSCloseCode.ABNORMAL_CLOSURE
365 self._exception = exc
366 self._response.close()
367 return True
369 if msg.type is WSMsgType.CLOSE:
370 self._close_code = msg.data
371 self._response.close()
372 return True
373 finally:
374 # Once closed the response can no longer be drained; release the
375 # parser and the stash it retains.
376 self._parser = None
378 @overload
379 async def receive(
380 self: "ClientWebSocketResponse[Literal[True]]", timeout: float | None = None
381 ) -> WSMessageDecodeText: ...
383 @overload
384 async def receive(
385 self: "ClientWebSocketResponse[Literal[False]]", timeout: float | None = None
386 ) -> WSMessageNoDecodeText: ...
388 @overload
389 async def receive(
390 self: "ClientWebSocketResponse[_DecodeText]", timeout: float | None = None
391 ) -> WSMessageDecodeText | WSMessageNoDecodeText: ...
393 async def receive(
394 self, timeout: float | None = None
395 ) -> WSMessageDecodeText | WSMessageNoDecodeText:
396 receive_timeout = timeout or self._timeout.ws_receive
398 while True:
399 if self._waiting:
400 raise RuntimeError("Concurrent call to receive() is not allowed")
402 if self._closed:
403 return WS_CLOSED_MESSAGE
404 elif self._closing:
405 await self.close()
406 return WS_CLOSED_MESSAGE
408 try:
409 self._waiting = True
410 try:
411 if receive_timeout:
412 # Entering the context manager and creating
413 # Timeout() object can take almost 50% of the
414 # run time in this loop so we avoid it if
415 # there is no read timeout.
416 async with asyncio.timeout(receive_timeout):
417 msg = await self._reader.read()
418 else:
419 msg = await self._reader.read()
420 finally:
421 self._waiting = False
422 if self._close_wait:
423 set_result(self._close_wait, None)
424 except (asyncio.CancelledError, TimeoutError):
425 self._close_code = WSCloseCode.ABNORMAL_CLOSURE
426 raise
427 except EofStream:
428 self._close_code = WSCloseCode.OK
429 await self.close()
430 return WS_CLOSED_MESSAGE
431 except ClientError:
432 # Likely ServerDisconnectedError when connection is lost.
433 # close() is not called on this path, so release the parser
434 # and the stash it retains here.
435 self._parser = None
436 self._set_closed()
437 self._close_code = WSCloseCode.ABNORMAL_CLOSURE
438 return WS_CLOSED_MESSAGE
439 except WebSocketError as exc:
440 self._close_code = exc.code
441 await self.close(code=exc.code)
442 return WSMessageError(data=exc)
443 except Exception as exc:
444 self._exception = exc
445 self._set_closing()
446 self._close_code = WSCloseCode.ABNORMAL_CLOSURE
447 await self.close()
448 return WSMessageError(data=exc)
450 if msg.type not in _INTERNAL_RECEIVE_TYPES:
451 # If its not a close/closing/ping/pong message
452 # we can return it immediately
453 return msg
455 if msg.type is WSMsgType.CLOSE:
456 self._set_closing()
457 self._close_code = msg.data
458 # Could be closed elsewhere while awaiting reader
459 if not self._closed and self._autoclose: # type: ignore[redundant-expr]
460 await self.close()
461 elif msg.type is WSMsgType.CLOSING:
462 self._set_closing()
463 elif msg.type is WSMsgType.PING and self._autoping:
464 await self.pong(msg.data)
465 continue
466 elif msg.type is WSMsgType.PONG and self._autoping:
467 continue
469 return msg
471 @overload
472 async def receive_str(
473 self: "ClientWebSocketResponse[Literal[True]]", *, timeout: float | None = None
474 ) -> str: ...
476 @overload
477 async def receive_str(
478 self: "ClientWebSocketResponse[Literal[False]]", *, timeout: float | None = None
479 ) -> bytes: ...
481 @overload
482 async def receive_str(
483 self: "ClientWebSocketResponse[_DecodeText]", *, timeout: float | None = None
484 ) -> str | bytes: ...
486 async def receive_str(self, *, timeout: float | None = None) -> str | bytes:
487 """Receive TEXT message.
489 Returns str when decode_text=True (default), bytes when decode_text=False.
490 """
491 msg = await self.receive(timeout)
492 if msg.type is not WSMsgType.TEXT:
493 raise WSMessageTypeError(
494 f"Received message {msg.type}:{msg.data!r} is not WSMsgType.TEXT"
495 )
496 return msg.data
498 async def receive_bytes(self, *, timeout: float | None = None) -> bytes:
499 msg = await self.receive(timeout)
500 if msg.type is not WSMsgType.BINARY:
501 raise WSMessageTypeError(
502 f"Received message {msg.type}:{msg.data!r} is not WSMsgType.BINARY"
503 )
504 return msg.data
506 @overload
507 async def receive_json(
508 self: "ClientWebSocketResponse[Literal[True]]",
509 *,
510 loads: JSONDecoder = ...,
511 timeout: float | None = None,
512 ) -> Any: ...
514 @overload
515 async def receive_json(
516 self: "ClientWebSocketResponse[Literal[False]]",
517 *,
518 loads: Callable[[bytes], Any] = ...,
519 timeout: float | None = None,
520 ) -> Any: ...
522 @overload
523 async def receive_json(
524 self: "ClientWebSocketResponse[_DecodeText]",
525 *,
526 loads: JSONDecoder | Callable[[bytes], Any] = ...,
527 timeout: float | None = None,
528 ) -> Any: ...
530 async def receive_json(
531 self,
532 *,
533 loads: JSONDecoder | Callable[[bytes], Any] = DEFAULT_JSON_DECODER,
534 timeout: float | None = None,
535 ) -> Any:
536 data = await self.receive_str(timeout=timeout)
537 return loads(data) # type: ignore[arg-type]
539 def __aiter__(self) -> Self:
540 return self
542 @overload
543 async def __anext__(
544 self: "ClientWebSocketResponse[Literal[True]]",
545 ) -> WSMessageDecodeText: ...
547 @overload
548 async def __anext__(
549 self: "ClientWebSocketResponse[Literal[False]]",
550 ) -> WSMessageNoDecodeText: ...
552 @overload
553 async def __anext__(
554 self: "ClientWebSocketResponse[_DecodeText]",
555 ) -> WSMessageDecodeText | WSMessageNoDecodeText: ...
557 async def __anext__(self) -> WSMessageDecodeText | WSMessageNoDecodeText:
558 msg = await self.receive()
559 if msg.type in (WSMsgType.CLOSE, WSMsgType.CLOSING, WSMsgType.CLOSED):
560 raise StopAsyncIteration
561 return msg
563 async def __aenter__(self) -> Self:
564 return self
566 async def __aexit__(
567 self,
568 exc_type: type[BaseException] | None,
569 exc_val: BaseException | None,
570 exc_tb: TracebackType | None,
571 ) -> None:
572 await self.close()