datastructures.py 4.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165
  1. """
  2. :mod:`websockets.datastructures` defines a class for manipulating HTTP headers.
  3. """
  4. from typing import (
  5. Any,
  6. Dict,
  7. Iterable,
  8. Iterator,
  9. List,
  10. Mapping,
  11. MutableMapping,
  12. Tuple,
  13. Union,
  14. )
  15. __all__ = ["Headers", "HeadersLike", "MultipleValuesError"]
  16. class MultipleValuesError(LookupError):
  17. """
  18. Exception raised when :class:`Headers` has more than one value for a key.
  19. """
  20. def __str__(self) -> str:
  21. # Implement the same logic as KeyError_str in Objects/exceptions.c.
  22. if len(self.args) == 1:
  23. return repr(self.args[0])
  24. return super().__str__()
  25. class Headers(MutableMapping[str, str]):
  26. """
  27. Efficient data structure for manipulating HTTP headers.
  28. A :class:`list` of ``(name, values)`` is inefficient for lookups.
  29. A :class:`dict` doesn't suffice because header names are case-insensitive
  30. and multiple occurrences of headers with the same name are possible.
  31. :class:`Headers` stores HTTP headers in a hybrid data structure to provide
  32. efficient insertions and lookups while preserving the original data.
  33. In order to account for multiple values with minimal hassle,
  34. :class:`Headers` follows this logic:
  35. - When getting a header with ``headers[name]``:
  36. - if there's no value, :exc:`KeyError` is raised;
  37. - if there's exactly one value, it's returned;
  38. - if there's more than one value, :exc:`MultipleValuesError` is raised.
  39. - When setting a header with ``headers[name] = value``, the value is
  40. appended to the list of values for that header.
  41. - When deleting a header with ``del headers[name]``, all values for that
  42. header are removed (this is slow).
  43. Other methods for manipulating headers are consistent with this logic.
  44. As long as no header occurs multiple times, :class:`Headers` behaves like
  45. :class:`dict`, except keys are lower-cased to provide case-insensitivity.
  46. Two methods support manipulating multiple values explicitly:
  47. - :meth:`get_all` returns a list of all values for a header;
  48. - :meth:`raw_items` returns an iterator of ``(name, values)`` pairs.
  49. """
  50. __slots__ = ["_dict", "_list"]
  51. def __init__(self, *args: Any, **kwargs: str) -> None:
  52. self._dict: Dict[str, List[str]] = {}
  53. self._list: List[Tuple[str, str]] = []
  54. # MutableMapping.update calls __setitem__ for each (name, value) pair.
  55. self.update(*args, **kwargs)
  56. def __str__(self) -> str:
  57. return "".join(f"{key}: {value}\r\n" for key, value in self._list) + "\r\n"
  58. def __repr__(self) -> str:
  59. return f"{self.__class__.__name__}({self._list!r})"
  60. def copy(self) -> "Headers":
  61. copy = self.__class__()
  62. copy._dict = self._dict.copy()
  63. copy._list = self._list.copy()
  64. return copy
  65. def serialize(self) -> bytes:
  66. # Headers only contain ASCII characters.
  67. return str(self).encode()
  68. # Collection methods
  69. def __contains__(self, key: object) -> bool:
  70. return isinstance(key, str) and key.lower() in self._dict
  71. def __iter__(self) -> Iterator[str]:
  72. return iter(self._dict)
  73. def __len__(self) -> int:
  74. return len(self._dict)
  75. # MutableMapping methods
  76. def __getitem__(self, key: str) -> str:
  77. value = self._dict[key.lower()]
  78. if len(value) == 1:
  79. return value[0]
  80. else:
  81. raise MultipleValuesError(key)
  82. def __setitem__(self, key: str, value: str) -> None:
  83. self._dict.setdefault(key.lower(), []).append(value)
  84. self._list.append((key, value))
  85. def __delitem__(self, key: str) -> None:
  86. key_lower = key.lower()
  87. self._dict.__delitem__(key_lower)
  88. # This is inefficent. Fortunately deleting HTTP headers is uncommon.
  89. self._list = [(k, v) for k, v in self._list if k.lower() != key_lower]
  90. def __eq__(self, other: Any) -> bool:
  91. if not isinstance(other, Headers):
  92. return NotImplemented
  93. return self._list == other._list
  94. def clear(self) -> None:
  95. """
  96. Remove all headers.
  97. """
  98. self._dict = {}
  99. self._list = []
  100. # Methods for handling multiple values
  101. def get_all(self, key: str) -> List[str]:
  102. """
  103. Return the (possibly empty) list of all values for a header.
  104. :param key: header name
  105. """
  106. return self._dict.get(key.lower(), [])
  107. def raw_items(self) -> Iterator[Tuple[str, str]]:
  108. """
  109. Return an iterator of all values as ``(name, value)`` pairs.
  110. """
  111. return iter(self._list)
  112. HeadersLike = Union[Headers, Mapping[str, str], Iterable[Tuple[str, str]]]
  113. HeadersLike__doc__ = """Types accepted wherever :class:`Headers` is expected"""
  114. # Remove try / except when dropping support for Python < 3.7
  115. try:
  116. HeadersLike.__doc__ = HeadersLike__doc__
  117. except AttributeError: # pragma: no cover
  118. pass