codekingpro/portable-devtools

1from __future__ import annotations
2
3from collections.abc import Iterable, Iterator, Mapping, MutableMapping
4from typing import Any, Protocol
5
6
7__all__ = [
8    "Headers",
9    "HeadersLike",
10    "MultipleValuesError",
11]
12
13
14class MultipleValuesError(LookupError):
15    """
16    Exception raised when :class:`Headers` has multiple values for a key.
17
18    """
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
26
27class Headers(MutableMapping[str, str]):
28    """
29    Efficient data structure for manipulating HTTP headers.
30
31    A :class:`list` of ``(name, values)`` is inefficient for lookups.
32
33    A :class:`dict` doesn't suffice because header names are case-insensitive
34    and multiple occurrences of headers with the same name are possible.
35
36    :class:`Headers` stores HTTP headers in a hybrid data structure to provide
37    efficient insertions and lookups while preserving the original data.
38
39    In order to account for multiple values with minimal hassle,
40    :class:`Headers` follows this logic:
41
42    - When getting a header with ``headers[name]``:
43        - if there's no value, :exc:`KeyError` is raised;
44        - if there's exactly one value, it's returned;
45        - if there's more than one value, :exc:`MultipleValuesError` is raised.
46
47    - When setting a header with ``headers[name] = value``, the value is
48      appended to the list of values for that header.
49
50    - When deleting a header with ``del headers[name]``, all values for that
51      header are removed (this is slow).
52
53    Other methods for manipulating headers are consistent with this logic.
54
55    As long as no header occurs multiple times, :class:`Headers` behaves like
56    :class:`dict`, except keys are lower-cased to provide case-insensitivity.
57
58    Two methods support manipulating multiple values explicitly:
59
60    - :meth:`get_all` returns a list of all values for a header;
61    - :meth:`raw_items` returns an iterator of ``(name, values)`` pairs.
62
63    """
64
65    __slots__ = ["_dict", "_list"]
66
67    # Like dict, Headers accepts an optional "mapping or iterable" argument.
68    def __init__(self, *args: HeadersLike, **kwargs: str) -> None:
69        self._dict: dict[str, list[str]] = {}
70        self._list: list[tuple[str, str]] = []
71        self.update(*args, **kwargs)
72
73    def __str__(self) -> str:
74        return "".join(f"{key}: {value}\r\n" for key, value in self._list) + "\r\n"
75
76    def __repr__(self) -> str:
77        return f"{self.__class__.__name__}({self._list!r})"
78
79    def copy(self) -> Headers:
80        copy = self.__class__()
81        copy._dict = self._dict.copy()
82        copy._list = self._list.copy()
83        return copy
84
85    def serialize(self) -> bytes:
86        # Since headers only contain ASCII characters, we can keep this simple.
87        return str(self).encode()
88
89    # Collection methods
90
91    def __contains__(self, key: object) -> bool:
92        return isinstance(key, str) and key.lower() in self._dict
93
94    def __iter__(self) -> Iterator[str]:
95        return iter(self._dict)
96
97    def __len__(self) -> int:
98        return len(self._dict)
99
100    # MutableMapping methods
101
102    def __getitem__(self, key: str) -> str:
103        value = self._dict[key.lower()]
104        if len(value) == 1:
105            return value[0]
106        else:
107            raise MultipleValuesError(key)
108
109    def __setitem__(self, key: str, value: str) -> None:
110        self._dict.setdefault(key.lower(), []).append(value)
111        self._list.append((key, value))
112
113    def __delitem__(self, key: str) -> None:
114        key_lower = key.lower()
115        self._dict.__delitem__(key_lower)
116        # This is inefficient. Fortunately deleting HTTP headers is uncommon.
117        self._list = [(k, v) for k, v in self._list if k.lower() != key_lower]
118
119    def __eq__(self, other: Any) -> bool:
120        if not isinstance(other, Headers):
121            return NotImplemented
122        return self._dict == other._dict
123
124    def clear(self) -> None:
125        """
126        Remove all headers.
127
128        """
129        self._dict = {}
130        self._list = []
131
132    def update(self, *args: HeadersLike, **kwargs: str) -> None:
133        """
134        Update from a :class:`Headers` instance and/or keyword arguments.
135
136        """
137        args = tuple(
138            arg.raw_items() if isinstance(arg, Headers) else arg for arg in args
139        )
140        super().update(*args, **kwargs)
141
142    # Methods for handling multiple values
143
144    def get_all(self, key: str) -> list[str]:
145        """
146        Return the (possibly empty) list of all values for a header.
147
148        Args:
149            key: Header name.
150
151        """
152        return self._dict.get(key.lower(), [])
153
154    def raw_items(self) -> Iterator[tuple[str, str]]:
155        """
156        Return an iterator of all values as ``(name, value)`` pairs.
157
158        """
159        return iter(self._list)
160
161
162# copy of _typeshed.SupportsKeysAndGetItem.
163class SupportsKeysAndGetItem(Protocol):  # pragma: no cover
164    """
165    Dict-like types with ``keys() -> str`` and ``__getitem__(key: str) -> str`` methods.
166
167    """
168
169    def keys(self) -> Iterable[str]: ...
170
171    def __getitem__(self, key: str) -> str: ...
172
173
174HeadersLike = (
175    Headers | Mapping[str, str] | Iterable[tuple[str, str]] | SupportsKeysAndGetItem
176)
177"""
178Types accepted where :class:`Headers` is expected.
179
180In addition to :class:`Headers` itself, this includes dict-like types where both
181keys and values are :class:`str`.
182
183"""
184