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