Coverage for /pythoncovmergedfiles/medio/medio/src/aiohttp/aiohttp/_websocket/models.py: 94%

1"""Models for WebSocket protocol versions 13 and 8.""" 

2 

3import json 

4from collections.abc import Callable 

5from enum import IntEnum 

6from typing import Any, Final, Literal, NamedTuple, cast 

7 

8WS_DEFLATE_TRAILING: Final[bytes] = bytes([0x00, 0x00, 0xFF, 0xFF]) 

9 

10 

11class WSCloseCode(IntEnum): 

12 OK = 1000 

13 GOING_AWAY = 1001 

14 PROTOCOL_ERROR = 1002 

15 UNSUPPORTED_DATA = 1003 

16 ABNORMAL_CLOSURE = 1006 

17 INVALID_TEXT = 1007 

18 POLICY_VIOLATION = 1008 

19 MESSAGE_TOO_BIG = 1009 

20 MANDATORY_EXTENSION = 1010 

21 INTERNAL_ERROR = 1011 

22 SERVICE_RESTART = 1012 

23 TRY_AGAIN_LATER = 1013 

24 BAD_GATEWAY = 1014 

25 

26 

27class WSMsgType(IntEnum): 

28 # websocket spec types 

29 CONTINUATION = 0x0 

30 TEXT = 0x1 

31 BINARY = 0x2 

32 PING = 0x9 

33 PONG = 0xA 

34 CLOSE = 0x8 

35 

36 # aiohttp specific types 

37 CLOSING = 0x100 

38 CLOSED = 0x101 

39 ERROR = 0x102 

40 

41 

42class WSMessageContinuation(NamedTuple): 

43 data: bytes 

44 size: int 

45 extra: str | None = None 

46 type: Literal[WSMsgType.CONTINUATION] = WSMsgType.CONTINUATION 

47 

48 

49class WSMessageText(NamedTuple): 

50 data: str 

51 size: int 

52 extra: str | None = None 

53 type: Literal[WSMsgType.TEXT] = WSMsgType.TEXT 

54 

55 def json( 

56 self, *, loads: Callable[[str | bytes | bytearray], Any] = json.loads 

57 ) -> Any: 

58 """Return parsed JSON data.""" 

59 return loads(self.data) 

60 

61 

62class WSMessageTextBytes(NamedTuple): 

63 """WebSocket TEXT message with raw bytes (no UTF-8 decoding).""" 

64 

65 data: bytes 

66 size: int 

67 extra: str | None = None 

68 type: Literal[WSMsgType.TEXT] = WSMsgType.TEXT 

69 

70 def json(self, *, loads: Callable[[bytes], Any] = json.loads) -> Any: 

71 """Return parsed JSON data.""" 

72 return loads(self.data) 

73 

74 

75class WSMessageBinary(NamedTuple): 

76 data: bytes 

77 size: int 

78 extra: str | None = None 

79 type: Literal[WSMsgType.BINARY] = WSMsgType.BINARY 

80 

81 def json( 

82 self, *, loads: Callable[[str | bytes | bytearray], Any] = json.loads 

83 ) -> Any: 

84 """Return parsed JSON data.""" 

85 return loads(self.data) 

86 

87 

88class WSMessagePing(NamedTuple): 

89 data: bytes 

90 size: int 

91 extra: str | None = None 

92 type: Literal[WSMsgType.PING] = WSMsgType.PING 

93 

94 

95class WSMessagePong(NamedTuple): 

96 data: bytes 

97 size: int 

98 extra: str | None = None 

99 type: Literal[WSMsgType.PONG] = WSMsgType.PONG 

100 

101 

102class WSMessageClose(NamedTuple): 

103 data: int 

104 size: int 

105 extra: str | None = None 

106 type: Literal[WSMsgType.CLOSE] = WSMsgType.CLOSE 

107 

108 

109class WSMessageClosing(NamedTuple): 

110 data: None = None 

111 size: int = 0 

112 extra: str | None = None 

113 type: Literal[WSMsgType.CLOSING] = WSMsgType.CLOSING 

114 

115 

116class WSMessageClosed(NamedTuple): 

117 data: None = None 

118 size: int = 0 

119 extra: str | None = None 

120 type: Literal[WSMsgType.CLOSED] = WSMsgType.CLOSED 

121 

122 

123class WSMessageError(NamedTuple): 

124 data: BaseException 

125 size: int = 0 

126 extra: str | None = None 

127 type: Literal[WSMsgType.ERROR] = WSMsgType.ERROR 

128 

129 

130# Base message types (excluding TEXT variants) 

131_WSMessageBase = ( 

132 WSMessageContinuation 

133 | WSMessageBinary 

134 | WSMessagePing 

135 | WSMessagePong 

136 | WSMessageClose 

137 | WSMessageClosing 

138 | WSMessageClosed 

139 | WSMessageError 

140) 

141 

142# All message types 

143WSMessage = _WSMessageBase | WSMessageText | WSMessageTextBytes 

144 

145# Message type when decode_text=True (default) - TEXT messages have str data 

146WSMessageDecodeText = _WSMessageBase | WSMessageText 

147 

148# Message type when decode_text=False - TEXT messages have bytes data 

149WSMessageNoDecodeText = _WSMessageBase | WSMessageTextBytes 

150 

151WS_CLOSED_MESSAGE = WSMessageClosed() 

152WS_CLOSING_MESSAGE = WSMessageClosing() 

153 

154 

155class WebSocketError(Exception): 

156 """WebSocket protocol parser error.""" 

157 

158 def __init__(self, code: int, message: str) -> None: 

159 self.code = code 

160 super().__init__(code, message) 

161 

162 def __str__(self) -> str: 

163 return cast(str, self.args[1]) 

164 

165 

166class WSHandshakeError(Exception): 

167 """WebSocket protocol handshake error."""