Move to bytestrings everywhere for headers.

Author: Lukasa

hyper/common/headers.py

@@ -7,6 +7,8 @@
 """
 import collections
 
+from hyper.compat import unicode, bytes, imap
+
 
 class HTTPHeaderMap(collections.MutableMapping):
     """
@@ -24,6 +26,24 @@ class HTTPHeaderMap(collections.MutableMapping):
 
     This data structure is an attempt to preserve all of that information
     while being as user-friendly as possible.
+
+    When iterated over, this structure returns headers in 'canonical form'.
+    This form is a tuple, where the first entry is the header name (in
+    lower-case), and the second entry is a list of header values (in original
+    case).
+
+    The mapping always emits both names and values in the form of bytestrings:
+    never unicode strings. It can accept names and values in unicode form, and
+    will automatically be encoded to bytestrings using UTF-8. The reason for
+    what appears to be a user-unfriendly decision here is primarily to allow
+    the broadest-possible compatibility (to make it possible to send headers in
+    unusual encodings) while ensuring that users are never confused about what
+    type of data they will receive.
+
+    ..warning:: Note that this data structure makes none of the performance
+                guarantees of a dictionary. Lookup and deletion is not an O(1)
+                operation. Inserting a new value *is* O(1), all other
+                operations are O(n), including *replacing* a header entirely.
     """
     def __init__(self, *args, **kwargs):
         # The meat of the structure. In practice, headers are an ordered list
@@ -42,17 +62,18 @@ def __init__(self, *args, **kwargs):
         self._items = []
 
         for arg in args:
-            self._items.extend(arg)
+            self._items.extend(map(lambda x: _to_bytestring_tuple(*x), arg))
 
         for k, v in kwargs.items():
-            self._items.append((k, v))
+            self._items.append(_to_bytestring_tuple(k, v))
 
     def __getitem__(self, key):
         """
         Unlike the dict __getitem__, this returns a list of items in the order
         they were added. These items are returned in 'canonical form', meaning
         that comma-separated values are split into multiple values.
         """
+        key = _to_bytestring(key)
         values = []
 
         for k, v in self._items:
@@ -68,14 +89,15 @@ def __setitem__(self, key, value):
         """
         Unlike the dict __setitem__, this appends to the list of items.
         """
-        self._items.append((key, value))
+        self._items.append(_to_bytestring_tuple(key, value))
 
     def __delitem__(self, key):
         """
         Sadly, __delitem__ is kind of stupid here, but the best we can do is
         delete all headers with a given key. To correctly achieve the 'KeyError
         on missing key' logic from dictionaries, we need to do this slowly.
         """
+        key = _to_bytestring(key)
         indices = []
         for (i, (k, v)) in enumerate(self._items):
             if _keys_equal(k, key):
@@ -111,6 +133,7 @@ def __contains__(self, key):
         """
         If any header is present with this key, returns True.
         """
+        key = _to_bytestring(key)
         return any(_keys_equal(key, k) for k, _ in self._items)
 
     def keys(self):
@@ -169,16 +192,37 @@ def canonical_form(k, v):
     canonical form. This means that the header is split on commas unless for
     any reason it's a super-special snowflake (I'm looking at you Set-Cookie).
     """
-    SPECIAL_SNOWFLAKES = set(['set-cookie', 'set-cookie2'])
+    SPECIAL_SNOWFLAKES = set([b'set-cookie', b'set-cookie2'])
 
     k = k.lower()
 
     if k in SPECIAL_SNOWFLAKES:
         yield k, v
     else:
-        for sub_val in v.split(','):
+        for sub_val in v.split(b','):
             yield k, sub_val.strip()
 
+
+def _to_bytestring(element):
+    """
+    Converts a single string to a bytestring, encoding via UTF-8 if needed.
+    """
+    if isinstance(element, unicode):
+        return element.encode('utf-8')
+    elif isinstance(element, bytes):
+        return element
+    else:
+        raise ValueError("Non string type.")
+
+
+def _to_bytestring_tuple(*x):
+    """
+    Converts the given strings to a bytestring if necessary, returning a
+    tuple.
+    """
+    return tuple(imap(_to_bytestring, x))
+
+
 def _keys_equal(x, y):
     """
     Returns 'True' if the two keys are equal by the laws of HTTP headers.

hyper/compat.py

@@ -36,6 +36,7 @@ def ignore_missing():
 
     from urllib import urlencode
     from urlparse import urlparse, urlsplit
+    from itertools import imap
 
     def to_byte(char):
         return ord(char)
@@ -52,9 +53,14 @@ def zlib_compressobj(level=6, method=zlib.DEFLATED, wbits=15, memlevel=8,
                          strategy=zlib.Z_DEFAULT_STRATEGY):
         return zlib.compressobj(level, method, wbits, memlevel, strategy)
 
+    unicode = unicode
+    bytes = str
+
 elif is_py3:
     from urllib.parse import urlencode, urlparse, urlsplit
 
+    imap = map
+
     def to_byte(char):
         return char
 
@@ -71,3 +77,6 @@ def write_to_stdout(data):
         ssl = ssl_compat
     else:
         import ssl
+
+    unicode = str
+    bytes = bytes

test/test_headers.py

@@ -6,19 +6,19 @@ class TestHTTPHeaderMap(object):
     def test_header_map_can_insert_single_header(self):
         h = HTTPHeaderMap()
         h['key'] = 'value'
-        assert h['key'] == ['value']
+        assert h['key'] == [b'value']
 
     def test_header_map_insensitive_key(self):
         h = HTTPHeaderMap()
         h['KEY'] = 'value'
-        assert h['key'] == ['value']
+        assert h['key'] == [b'value']
 
     def test_header_map_is_iterable_in_order(self):
         h = HTTPHeaderMap()
         items = [
-            ('k1', 'v2'),
-            ('k2', 'v2'),
-            ('k2', 'v3'),
+            (b'k1', b'v2'),
+            (b'k2', b'v2'),
+            (b'k2', b'v3'),
         ]
 
         for k, v in items:
@@ -29,18 +29,18 @@ def test_header_map_is_iterable_in_order(self):
 
     def test_header_map_allows_multiple_values(self):
         h = HTTPHeaderMap()
-        h['key'] = 'v1'
-        h['Key'] = 'v2'
+        h['key'] = b'v1'
+        h[b'Key'] = b'v2'
 
-        assert h['key'] == ['v1', 'v2']
+        assert h['key'] == [b'v1', b'v2']
 
     def test_header_map_can_delete_value(self):
         h = HTTPHeaderMap()
-        h['key'] = 'v1'
-        del h['key']
+        h['key'] = b'v1'
+        del h[b'key']
 
         with pytest.raises(KeyError):
-            h['key']
+            h[b'key']
 
     def test_header_map_deletes_all_values(self):
         h = HTTPHeaderMap()
@@ -55,13 +55,14 @@ def test_setting_comma_separated_header(self):
         h = HTTPHeaderMap()
         h['key'] = 'v1, v2'
 
-        assert h['key'] == ['v1', 'v2']
+        assert h[b'key'] == [b'v1', b'v2']
 
     def test_containment(self):
         h = HTTPHeaderMap()
         h['key'] = 'val'
 
         assert 'key' in h
+        assert b'key' in h
         assert 'nonkey' not in h
 
     def test_length_counts_lines_separately(self):
@@ -79,7 +80,7 @@ def test_keys(self):
         h['k1'] = 'v4'
 
         assert len(list(h.keys())) == 4
-        assert list(h.keys()) == ['k1', 'k1', 'k2', 'k1']
+        assert list(h.keys()) == [b'k1', b'k1', b'k2', b'k1']
 
     def test_values(self):
         h = HTTPHeaderMap()
@@ -88,14 +89,14 @@ def test_values(self):
         h['k1'] = 'v4'
 
         assert len(list(h.values())) == 4
-        assert list(h.values()) == ['v1', 'v2', 'v3', 'v4']
+        assert list(h.values()) == [b'v1', b'v2', b'v3', b'v4']
 
     def test_items(self):
         h = HTTPHeaderMap()
         items = [
-            ('k1', 'v2'),
-            ('k2', 'v2'),
-            ('k2', 'v3'),
+            (b'k1', b'v2'),
+            (b'k2', b'v2'),
+            (b'k2', b'v3'),
         ]
 
         for k, v in items:
@@ -114,13 +115,13 @@ def test_actual_get(self):
         h['k2'] = 'v3'
         h['k1'] = 'v4'
 
-        assert h.get('k1') == ['v1', 'v2', 'v4']
+        assert h.get('k1') == [b'v1', b'v2', b'v4']
 
     def test_doesnt_split_set_cookie(self):
         h = HTTPHeaderMap()
         h['Set-Cookie'] = 'v1, v2'
-        assert h['set-cookie'] == ['v1, v2']
-        assert h.get('set-cookie') == ['v1, v2']
+        assert h['set-cookie'] == [b'v1, v2']
+        assert h.get(b'set-cookie') == [b'v1, v2']
 
     def test_equality(self):
         h1 = HTTPHeaderMap()
@@ -171,44 +172,44 @@ def test_deleting_nonexistent(self):
 
     def test_can_create_from_iterable(self):
         items = [
-            ('k1', 'v2'),
-            ('k2', 'v2'),
-            ('k2', 'v3'),
+            (b'k1', b'v2'),
+            (b'k2', b'v2'),
+            (b'k2', b'v3'),
         ]
         h = HTTPHeaderMap(items)
 
         assert list(h) == items
 
     def test_can_create_from_multiple_iterables(self):
         items = [
-            ('k1', 'v2'),
-            ('k2', 'v2'),
-            ('k2', 'v3'),
+            (b'k1', b'v2'),
+            (b'k2', b'v2'),
+            (b'k2', b'v3'),
         ]
         h = HTTPHeaderMap(items, items, items)
 
         assert list(h) == items + items + items
 
     def test_create_from_iterables_and_kwargs(self):
         items = [
-            ('k1', 'v2'),
-            ('k2', 'v2'),
-            ('k2', 'v3'),
+            (b'k1', b'v2'),
+            (b'k2', b'v2'),
+            (b'k2', b'v3'),
         ]
         h = list(HTTPHeaderMap(items, k3='v4', k4='v5'))
 
         # kwargs are an unordered dictionary, so allow for both possible
         # iteration orders.
         assert (
-            h == items + [('k3', 'v4'), ('k4', 'v5')] or
-            h == items + [('k4', 'v5'), ('k3', 'v4')]
+            h == items + [(b'k3', b'v4'), (b'k4', b'v5')] or
+            h == items + [(b'k4', b'v5'), (b'k3', b'v4')]
         )
 
     def test_raw_iteration(self):
         items = [
-            ('k1', 'v2'),
-            ('k2', 'v2, v3, v4'),
-            ('k2', 'v3'),
+            (b'k1', b'v2'),
+            (b'k2', b'v2, v3, v4'),
+            (b'k2', b'v3'),
         ]
         h = HTTPHeaderMap(items)