master
py 320 lines 9.92 KB
Raw
1 # SPDX-License-Identifier: MIT
2 from __future__ import absolute_import
3
4 try:
5 from collections import Mapping, MutableMapping
6 except ImportError:
7 from collections.abc import Mapping, MutableMapping
8
9 try:
10 from threading import RLock
11 except ImportError: # Platform-specific: No threads available
12 class RLock:
13 def __enter__(self):
14 pass
15
16 def __exit__(self, exc_type, exc_value, traceback):
17 pass
18
19
20 try: # Python 2.7+
21 from collections import OrderedDict
22 except ImportError:
23 from .packages.ordered_dict import OrderedDict
24 from .packages.six import iterkeys, itervalues, PY3
25
26
27 __all__ = ['RecentlyUsedContainer', 'HTTPHeaderDict']
28
29
30 _Null = object()
31
32
33 class RecentlyUsedContainer(MutableMapping):
34 """
35 Provides a thread-safe dict-like container which maintains up to
36 ``maxsize`` keys while throwing away the least-recently-used keys beyond
37 ``maxsize``.
38
39 :param maxsize:
40 Maximum number of recent elements to retain.
41
42 :param dispose_func:
43 Every time an item is evicted from the container,
44 ``dispose_func(value)`` is called. Callback which will get called
45 """
46
47 ContainerCls = OrderedDict
48
49 def __init__(self, maxsize=10, dispose_func=None):
50 self._maxsize = maxsize
51 self.dispose_func = dispose_func
52
53 self._container = self.ContainerCls()
54 self.lock = RLock()
55
56 def __getitem__(self, key):
57 # Re-insert the item, moving it to the end of the eviction line.
58 with self.lock:
59 item = self._container.pop(key)
60 self._container[key] = item
61 return item
62
63 def __setitem__(self, key, value):
64 evicted_value = _Null
65 with self.lock:
66 # Possibly evict the existing value of 'key'
67 evicted_value = self._container.get(key, _Null)
68 self._container[key] = value
69
70 # If we didn't evict an existing value, we might have to evict the
71 # least recently used item from the beginning of the container.
72 if len(self._container) > self._maxsize:
73 _key, evicted_value = self._container.popitem(last=False)
74
75 if self.dispose_func and evicted_value is not _Null:
76 self.dispose_func(evicted_value)
77
78 def __delitem__(self, key):
79 with self.lock:
80 value = self._container.pop(key)
81
82 if self.dispose_func:
83 self.dispose_func(value)
84
85 def __len__(self):
86 with self.lock:
87 return len(self._container)
88
89 def __iter__(self):
90 raise NotImplementedError('Iteration over this class is unlikely to be threadsafe.')
91
92 def clear(self):
93 with self.lock:
94 # Copy pointers to all values, then wipe the mapping
95 values = list(itervalues(self._container))
96 self._container.clear()
97
98 if self.dispose_func:
99 for value in values:
100 self.dispose_func(value)
101
102 def keys(self):
103 with self.lock:
104 return list(iterkeys(self._container))
105
106
107 class HTTPHeaderDict(MutableMapping):
108 """
109 :param headers:
110 An iterable of field-value pairs. Must not contain multiple field names
111 when compared case-insensitively.
112
113 :param kwargs:
114 Additional field-value pairs to pass in to ``dict.update``.
115
116 A ``dict`` like container for storing HTTP Headers.
117
118 Field names are stored and compared case-insensitively in compliance with
119 RFC 7230. Iteration provides the first case-sensitive key seen for each
120 case-insensitive pair.
121
122 Using ``__setitem__`` syntax overwrites fields that compare equal
123 case-insensitively in order to maintain ``dict``'s api. For fields that
124 compare equal, instead create a new ``HTTPHeaderDict`` and use ``.add``
125 in a loop.
126
127 If multiple fields that are equal case-insensitively are passed to the
128 constructor or ``.update``, the behavior is undefined and some will be
129 lost.
130
131 >>> headers = HTTPHeaderDict()
132 >>> headers.add('Set-Cookie', 'foo=bar')
133 >>> headers.add('set-cookie', 'baz=quxx')
134 >>> headers['content-length'] = '7'
135 >>> headers['SET-cookie']
136 'foo=bar, baz=quxx'
137 >>> headers['Content-Length']
138 '7'
139 """
140
141 def __init__(self, headers=None, **kwargs):
142 super(HTTPHeaderDict, self).__init__()
143 self._container = OrderedDict()
144 if headers is not None:
145 if isinstance(headers, HTTPHeaderDict):
146 self._copy_from(headers)
147 else:
148 self.extend(headers)
149 if kwargs:
150 self.extend(kwargs)
151
152 def __setitem__(self, key, val):
153 self._container[key.lower()] = [key, val]
154 return self._container[key.lower()]
155
156 def __getitem__(self, key):
157 val = self._container[key.lower()]
158 return ', '.join(val[1:])
159
160 def __delitem__(self, key):
161 del self._container[key.lower()]
162
163 def __contains__(self, key):
164 return key.lower() in self._container
165
166 def __eq__(self, other):
167 if not isinstance(other, Mapping) and not hasattr(other, 'keys'):
168 return False
169 if not isinstance(other, type(self)):
170 other = type(self)(other)
171 return (dict((k.lower(), v) for k, v in self.itermerged()) ==
172 dict((k.lower(), v) for k, v in other.itermerged()))
173
174 def __ne__(self, other):
175 return not self.__eq__(other)
176
177 if not PY3: # Python 2
178 iterkeys = MutableMapping.iterkeys
179 itervalues = MutableMapping.itervalues
180
181 __marker = object()
182
183 def __len__(self):
184 return len(self._container)
185
186 def __iter__(self):
187 # Only provide the originally cased names
188 for vals in self._container.values():
189 yield vals[0]
190
191 def pop(self, key, default=__marker):
192 '''D.pop(k[,d]) -> v, remove specified key and return the corresponding value.
193 If key is not found, d is returned if given, otherwise KeyError is raised.
194 '''
195 # Using the MutableMapping function directly fails due to the private marker.
196 # Using ordinary dict.pop would expose the internal structures.
197 # So let's reinvent the wheel.
198 try:
199 value = self[key]
200 except KeyError:
201 if default is self.__marker:
202 raise
203 return default
204 else:
205 del self[key]
206 return value
207
208 def discard(self, key):
209 try:
210 del self[key]
211 except KeyError:
212 pass
213
214 def add(self, key, val):
215 """Adds a (name, value) pair, doesn't overwrite the value if it already
216 exists.
217
218 >>> headers = HTTPHeaderDict(foo='bar')
219 >>> headers.add('Foo', 'baz')
220 >>> headers['foo']
221 'bar, baz'
222 """
223 key_lower = key.lower()
224 new_vals = [key, val]
225 # Keep the common case aka no item present as fast as possible
226 vals = self._container.setdefault(key_lower, new_vals)
227 if new_vals is not vals:
228 vals.append(val)
229
230 def extend(self, *args, **kwargs):
231 """Generic import function for any type of header-like object.
232 Adapted version of MutableMapping.update in order to insert items
233 with self.add instead of self.__setitem__
234 """
235 if len(args) > 1:
236 raise TypeError("extend() takes at most 1 positional "
237 "arguments ({0} given)".format(len(args)))
238 other = args[0] if len(args) >= 1 else ()
239
240 if isinstance(other, HTTPHeaderDict):
241 for key, val in other.iteritems():
242 self.add(key, val)
243 elif isinstance(other, Mapping):
244 for key in other:
245 self.add(key, other[key])
246 elif hasattr(other, "keys"):
247 for key in other.keys():
248 self.add(key, other[key])
249 else:
250 for key, value in other:
251 self.add(key, value)
252
253 for key, value in kwargs.items():
254 self.add(key, value)
255
256 def getlist(self, key):
257 """Returns a list of all the values for the named field. Returns an
258 empty list if the key doesn't exist."""
259 try:
260 vals = self._container[key.lower()]
261 except KeyError:
262 return []
263 else:
264 return vals[1:]
265
266 # Backwards compatibility for httplib
267 getheaders = getlist
268 getallmatchingheaders = getlist
269 iget = getlist
270
271 def __repr__(self):
272 return "%s(%s)" % (type(self).__name__, dict(self.itermerged()))
273
274 def _copy_from(self, other):
275 for key in other:
276 val = other.getlist(key)
277 if isinstance(val, list):
278 # Don't need to convert tuples
279 val = list(val)
280 self._container[key.lower()] = [key] + val
281
282 def copy(self):
283 clone = type(self)()
284 clone._copy_from(self)
285 return clone
286
287 def iteritems(self):
288 """Iterate over all header lines, including duplicate ones."""
289 for key in self:
290 vals = self._container[key.lower()]
291 for val in vals[1:]:
292 yield vals[0], val
293
294 def itermerged(self):
295 """Iterate over all headers, merging duplicate ones together."""
296 for key in self:
297 val = self._container[key.lower()]
298 yield val[0], ', '.join(val[1:])
299
300 def items(self):
301 return list(self.iteritems())
302
303 @classmethod
304 def from_httplib(cls, message): # Python 2
305 """Read headers from a Python 2 httplib message object."""
306 # python2.7 does not expose a proper API for exporting multiheaders
307 # efficiently. This function re-reads raw lines from the message
308 # object and extracts the multiheaders properly.
309 headers = []
310
311 for line in message.headers:
312 if line.startswith((' ', '\t')):
313 key, value = headers[-1]
314 headers[-1] = (key, value + '\r\n' + line.rstrip())
315 continue
316
317 key, value = line.split(':', 1)
318 headers.append((key, value.strip()))
319
320 return cls(headers)