1"""Models for WebSocket protocol versions 13 and 8."""
2
3import json
4from enum import IntEnum
5from typing import Any, Callable, Final, NamedTuple, Optional, cast
6
7WS_DEFLATE_TRAILING: Final[bytes] = bytes([0x00, 0x00, 0xFF, 0xFF])
8
9
10class WSCloseCode(IntEnum):
11 OK = 1000
12 GOING_AWAY = 1001
13 PROTOCOL_ERROR = 1002
14 UNSUPPORTED_DATA = 1003
15 ABNORMAL_CLOSURE = 1006
16 INVALID_TEXT = 1007
17 POLICY_VIOLATION = 1008
18 MESSAGE_TOO_BIG = 1009
19 MANDATORY_EXTENSION = 1010
20 INTERNAL_ERROR = 1011
21 SERVICE_RESTART = 1012
22 TRY_AGAIN_LATER = 1013
23 BAD_GATEWAY = 1014
24
25
26class WSMsgType(IntEnum):
27 # websocket spec types
28 CONTINUATION = 0x0
29 TEXT = 0x1
30 BINARY = 0x2
31 PING = 0x9
32 PONG = 0xA
33 CLOSE = 0x8
34
35 # aiohttp specific types
36 CLOSING = 0x100
37 CLOSED = 0x101
38 ERROR = 0x102
39
40 text = TEXT
41 binary = BINARY
42 ping = PING
43 pong = PONG
44 close = CLOSE
45 closing = CLOSING
46 closed = CLOSED
47 error = ERROR
48
49
50class WSMessage(NamedTuple):
51 type: WSMsgType
52 # To type correctly, this would need some kind of tagged union for each type.
53 data: Any
54 extra: Optional[str]
55
56 def json(self, *, loads: Callable[[Any], Any] = json.loads) -> Any:
57 """Return parsed JSON data.
58
59 .. versionadded:: 0.22
60 """
61 return loads(self.data)
62
63
64# Constructing the tuple directly to avoid the overhead of
65# the lambda and arg processing since NamedTuples are constructed
66# with a run time built lambda
67# https://github.com/python/cpython/blob/d83fcf8371f2f33c7797bc8f5423a8bca8c46e5c/Lib/collections/__init__.py#L441
68WS_CLOSED_MESSAGE = tuple.__new__(WSMessage, (WSMsgType.CLOSED, None, None))
69WS_CLOSING_MESSAGE = tuple.__new__(WSMessage, (WSMsgType.CLOSING, None, None))
70
71
72class WebSocketError(Exception):
73 """WebSocket protocol parser error."""
74
75 def __init__(self, code: int, message: str) -> None:
76 self.code = code
77 super().__init__(code, message)
78
79 def __str__(self) -> str:
80 return cast(str, self.args[1])
81
82
83class WSHandshakeError(Exception):
84 """WebSocket protocol handshake error."""