models.py 2.1 KB

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