datastructures.py 7.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225
  1. from __future__ import annotations
  2. import re
  3. from collections.abc import Iterable, Iterator, Mapping, MutableMapping
  4. from typing import Any, Protocol
  5. __all__ = [
  6. "Headers",
  7. "HeadersLike",
  8. "MultipleValuesError",
  9. ]
  10. class MultipleValuesError(LookupError):
  11. """
  12. Exception raised when :class:`Headers` has multiple values for a key.
  13. """
  14. def __str__(self) -> str:
  15. # Implement the same logic as KeyError_str in Objects/exceptions.c.
  16. if len(self.args) == 1:
  17. return repr(self.args[0])
  18. return super().__str__()
  19. # Same regex as http11._value_re, but for matching str rather than bytes.
  20. is_valid_header_value = re.compile(r"[\x09\x20-\x7e\x80-\xff]*").fullmatch
  21. class Headers(MutableMapping[str, str]):
  22. """
  23. Efficient data structure for manipulating HTTP headers.
  24. A :class:`list` of ``(name, values)`` is inefficient for lookups.
  25. A :class:`dict` doesn't suffice because header names are case-insensitive
  26. and multiple occurrences of headers with the same name are possible.
  27. :class:`Headers` stores HTTP headers in a hybrid data structure to provide
  28. efficient insertions and lookups while preserving the original data.
  29. In order to account for multiple values with minimal hassle,
  30. :class:`Headers` follows this logic:
  31. - When getting a header with ``headers[name]``:
  32. - if there's no value, :exc:`KeyError` is raised;
  33. - if there's exactly one value, it's returned;
  34. - if there's more than one value, :exc:`MultipleValuesError` is raised.
  35. - When setting a header with ``headers[name] = value``, the value is
  36. appended to the list of values for that header.
  37. - When deleting a header with ``del headers[name]``, all values for that
  38. header are removed (this is slow).
  39. Other methods for manipulating headers are consistent with this logic.
  40. As long as no header occurs multiple times, :class:`Headers` behaves like
  41. :class:`dict`, except keys are lower-cased to provide case-insensitivity.
  42. Two methods support manipulating multiple values explicitly:
  43. - :meth:`get_all` returns a list of all values for a header;
  44. - :meth:`raw_items` returns an iterator of ``(name, values)`` pairs.
  45. Header names and values are expected to contain only ASCII text. However,
  46. non-ASCII values happen in practice, even though there is no standard for
  47. transmitting non-ASCII data in HTTP headers. :class:`Headers` supports it
  48. by treating it as ISO-8859-1 data. This is a safe and reversible encoding
  49. to represent arbitrary data in a :class:`str`.
  50. When reading headers from the network, if the actual encoding isn't
  51. ISO-8859-1, you must re-encode and decode, e.g.::
  52. value = headers[key].encode("iso-8859-1").decode("utf-8")
  53. Conversely, when sending headers to the network, if you need to use a
  54. different encoding, you can encode and decode, e.g.::
  55. headers[key] = value.encode("utf-8").decode("iso-8859-1")
  56. When assigning a value to a header, as a security hardening measure, the
  57. value is checked for unsafe characters. The name isn't checked because it's
  58. usually a constant in code, unlikely to be tainted by user input.
  59. """
  60. __slots__ = ["_dict", "_list"]
  61. # Like dict, Headers accepts an optional "mapping or iterable" argument.
  62. def __init__(self, *args: HeadersLike, **kwargs: str) -> None:
  63. self._dict: dict[str, list[str]] = {}
  64. self._list: list[tuple[str, str]] = []
  65. self.update(*args, **kwargs)
  66. def __str__(self) -> str:
  67. return "".join(f"{key}: {value}\r\n" for key, value in self._list) + "\r\n"
  68. def __repr__(self) -> str:
  69. return f"{self.__class__.__name__}({self._list!r})"
  70. def copy(self) -> Headers:
  71. copy = self.__class__()
  72. copy._dict = self._dict.copy()
  73. copy._list = self._list.copy()
  74. return copy
  75. def serialize(self) -> bytes:
  76. # parse_headers() supports non-ASCII header values. It decodes them as
  77. # ISO-8859-1. Encode back in ISO-8859-1 in order to round-trip cleanly.
  78. return str(self).encode("iso-8859-1")
  79. # Collection methods
  80. def __contains__(self, key: object) -> bool:
  81. return isinstance(key, str) and key.lower() in self._dict
  82. def __iter__(self) -> Iterator[str]:
  83. return iter(self._dict)
  84. def __len__(self) -> int:
  85. return len(self._dict)
  86. # MutableMapping methods
  87. def __getitem__(self, key: str) -> str:
  88. value = self._dict[key.lower()]
  89. if len(value) == 1:
  90. return value[0]
  91. else:
  92. raise MultipleValuesError(key)
  93. def __setitem__(self, key: str, value: str) -> None:
  94. if not is_valid_header_value(str(value)):
  95. raise InvalidHeaderValue(key, value)
  96. self._dict.setdefault(key.lower(), []).append(value)
  97. self._list.append((key, value))
  98. def __delitem__(self, key: str) -> None:
  99. key_lower = key.lower()
  100. self._dict.__delitem__(key_lower)
  101. # This is inefficient. Fortunately deleting HTTP headers is uncommon.
  102. self._list = [(k, v) for k, v in self._list if k.lower() != key_lower]
  103. def __eq__(self, other: Any) -> bool:
  104. if not isinstance(other, Headers):
  105. return NotImplemented
  106. return self._dict == other._dict
  107. def clear(self) -> None:
  108. """
  109. Remove all headers.
  110. """
  111. self._dict = {}
  112. self._list = []
  113. def update(self, *args: HeadersLike, **kwargs: str) -> None:
  114. """
  115. Update from a :class:`Headers` instance and/or keyword arguments.
  116. """
  117. args = tuple(
  118. arg.raw_items() if isinstance(arg, Headers) else arg for arg in args
  119. )
  120. super().update(*args, **kwargs)
  121. # Methods for handling multiple values
  122. def get_all(self, key: str) -> list[str]:
  123. """
  124. Return the (possibly empty) list of all values for a header.
  125. Args:
  126. key: Header name.
  127. """
  128. return self._dict.get(key.lower(), [])
  129. def raw_items(self) -> Iterator[tuple[str, str]]:
  130. """
  131. Return an iterator of all values as ``(name, value)`` pairs.
  132. """
  133. return iter(self._list)
  134. # Internal methods
  135. def set_insecure(self, key: str, value: str) -> None:
  136. """
  137. Set a header without validating its value.
  138. """
  139. self._dict.setdefault(key.lower(), []).append(value)
  140. self._list.append((key, value))
  141. # copy of _typeshed.SupportsKeysAndGetItem.
  142. class SupportsKeysAndGetItem(Protocol):
  143. """
  144. Dict-like types with ``keys() -> str`` and ``__getitem__(key: str) -> str`` methods.
  145. """
  146. def keys(self) -> Iterable[str]: ... # pragma: no branch
  147. def __getitem__(self, key: str) -> str: ... # pragma: no branch
  148. HeadersLike = (
  149. Headers | Mapping[str, str] | Iterable[tuple[str, str]] | SupportsKeysAndGetItem
  150. )
  151. """
  152. Types accepted where :class:`Headers` is expected.
  153. In addition to :class:`Headers` itself, this includes dict-like types where both
  154. keys and values are :class:`str`.
  155. """
  156. # At the bottom to break an import cycle.
  157. from .exceptions import InvalidHeaderValue # noqa: E402