Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/aiohttp/_websocket/models.py: 92%

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, 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 text = TEXT 

42 binary = BINARY 

43 ping = PING 

44 pong = PONG 

45 close = CLOSE 

46 closing = CLOSING 

47 closed = CLOSED 

48 error = ERROR 

49 

50 

51class WSMessage(NamedTuple): 

52 type: WSMsgType 

53 # To type correctly, this would need some kind of tagged union for each type. 

54 data: Any 

55 extra: str | None 

56 

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

58 """Return parsed JSON data. 

59 

60 .. versionadded:: 0.22 

61 """ 

62 return loads(self.data) 

63 

64 

65class WSMessageTextBytes(NamedTuple): 

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

67 

68 type: WSMsgType 

69 # To type correctly, this would need some kind of tagged union for each type. 

70 # In 4.0, we use a union of message types to properly type data, but in 3.x 

71 # we keep it as Any to avoid a breaking change. 

72 data: Any 

73 extra: str | None 

74 

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

76 """Return parsed JSON data.""" 

77 return loads(self.data) 

78 

79 

80# Type aliases for message types based on decode_text setting 

81# When decode_text=True, TEXT messages have str data (WSMessage) 

82# When decode_text=False, TEXT messages have bytes data (WSMessageTextBytes) 

83WSMessageDecodeText = WSMessage 

84WSMessageNoDecodeText = WSMessage | WSMessageTextBytes 

85 

86 

87# Constructing the tuple directly to avoid the overhead of 

88# the lambda and arg processing since NamedTuples are constructed 

89# with a run time built lambda 

90# https://github.com/python/cpython/blob/d83fcf8371f2f33c7797bc8f5423a8bca8c46e5c/Lib/collections/__init__.py#L441 

91WS_CLOSED_MESSAGE = tuple.__new__(WSMessage, (WSMsgType.CLOSED, None, None)) 

92WS_CLOSING_MESSAGE = tuple.__new__(WSMessage, (WSMsgType.CLOSING, None, None)) 

93 

94 

95class WebSocketError(Exception): 

96 """WebSocket protocol parser error.""" 

97 

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

99 self.code = code 

100 super().__init__(code, message) 

101 

102 def __str__(self) -> str: 

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

104 

105 

106class WSHandshakeError(Exception): 

107 """WebSocket protocol handshake error."""