Minimal Headers data structure

Author: lovelydinosaur

h11/_connection.py

@@ -534,7 +534,7 @@ def send_failed(self):
     def _clean_up_response_headers_for_sending(self, response):
         assert type(response) is Response
 
-        headers = list(response.headers)
+        headers = response.headers
         need_close = False
 
         # HEAD requests need some special handling: they always act like they
@@ -560,27 +560,27 @@ def _clean_up_response_headers_for_sending(self, response):
             # but the HTTP spec says that if our peer does this then we have
             # to fix it instead of erroring out, so we'll accord the user the
             # same respect).
-            set_comma_header(headers, b"content-length", [])
+            headers = set_comma_header(headers, b"content-length", [])
             if self.their_http_version is None or self.their_http_version < b"1.1":
                 # Either we never got a valid request and are sending back an
                 # error (their_http_version is None), so we assume the worst;
                 # or else we did get a valid HTTP/1.0 request, so we know that
                 # they don't understand chunked encoding.
-                set_comma_header(headers, b"transfer-encoding", [])
+                headers = set_comma_header(headers, b"transfer-encoding", [])
                 # This is actually redundant ATM, since currently we
                 # unconditionally disable keep-alive when talking to HTTP/1.0
                 # peers. But let's be defensive just in case we add
                 # Connection: keep-alive support later:
                 if self._request_method != b"HEAD":
                     need_close = True
             else:
-                set_comma_header(headers, b"transfer-encoding", ["chunked"])
+                headers = set_comma_header(headers, b"transfer-encoding", ["chunked"])
 
         if not self._cstate.keep_alive or need_close:
             # Make sure Connection: close is set
             connection = set(get_comma_header(headers, b"connection"))
             connection.discard(b"keep-alive")
             connection.add(b"close")
-            set_comma_header(headers, b"connection", sorted(connection))
+            headers = set_comma_header(headers, b"connection", sorted(connection))
 
         response.headers = headers

h11/_headers.py

@@ -62,6 +62,23 @@
 _field_value_re = re.compile(field_value.encode("ascii"))
 
 
+class Headers:
+    __slots__ = ['raw_items']
+
+    def __init__(self, raw_items):
+        self.raw_items = raw_items
+
+    def __iter__(self):
+        for _, name, value in self.raw_items:
+            yield name, value
+
+    def __bool__(self):
+        return bool(self.raw_items)
+
+    def __eq__(self, other):
+        return list(self) == list(other)
+
+
 def normalize_and_validate(headers, _parsed=False):
     new_headers = []
     saw_content_length = False
@@ -75,6 +92,7 @@ def normalize_and_validate(headers, _parsed=False):
             value = bytesify(value)
             validate(_field_name_re, name, "Illegal header name {!r}", name)
             validate(_field_value_re, value, "Illegal header value {!r}", value)
+        raw_name = name
         name = name.lower()
         if name == b"content-length":
             if saw_content_length:
@@ -99,8 +117,8 @@ def normalize_and_validate(headers, _parsed=False):
                     error_status_hint=501,
                 )
             saw_transfer_encoding = True
-        new_headers.append((name, value))
-    return new_headers
+        new_headers.append((raw_name, name, value))
+    return Headers(new_headers)
 
 
 def get_comma_header(headers, name):
@@ -140,7 +158,7 @@ def get_comma_header(headers, name):
     # "100-continue". Splitting on commas is harmless. Case insensitive.
     #
     out = []
-    for found_name, found_raw_value in headers:
+    for _, found_name, found_raw_value in headers.raw_items:
         if found_name == name:
             found_raw_value = found_raw_value.lower()
             for found_split_value in found_raw_value.split(b","):
@@ -153,12 +171,12 @@ def get_comma_header(headers, name):
 def set_comma_header(headers, name, new_values):
     # The header name `name` is expected to be lower-case bytes.
     new_headers = []
-    for found_name, found_raw_value in headers:
+    for found_raw_name, found_name, found_raw_value in headers.raw_items:
         if found_name != name:
-            new_headers.append((found_name, found_raw_value))
+            new_headers.append((found_raw_name, found_raw_value))
     for new_value in new_values:
         new_headers.append((name, new_value))
-    headers[:] = normalize_and_validate(new_headers)
+    return normalize_and_validate(new_headers)
 
 
 def has_expect_100_continue(request):

h11/_writers.py

@@ -38,10 +38,11 @@ def write_headers(headers, write):
     # "Since the Host field-value is critical information for handling a
     # request, a user agent SHOULD generate Host as the first header field
     # following the request-line." - RFC 7230
-    for name, value in headers:
+    raw_items = headers.raw_items
+    for raw_name, name, value in raw_items:
         if name == b"host":
             write(bytesmod(b"%s: %s\r\n", (name, value)))
-    for name, value in headers:
+    for raw_name, name, value in raw_items:
         if name != b"host":
             write(bytesmod(b"%s: %s\r\n", (name, value)))
     write(b"\r\n")

h11/tests/test_headers.py

@@ -83,7 +83,7 @@ def test_get_set_comma_header():
 
     assert get_comma_header(headers, b"connection") == [b"close", b"foo", b"bar"]
 
-    set_comma_header(headers, b"newthing", ["a", "b"])
+    headers = set_comma_header(headers, b"newthing", ["a", "b"])
 
     with pytest.raises(LocalProtocolError):
         set_comma_header(headers, b"newthing", ["  a", "b"])
@@ -96,7 +96,7 @@ def test_get_set_comma_header():
         (b"newthing", b"b"),
     ]
 
-    set_comma_header(headers, b"whatever", ["different thing"])
+    headers = set_comma_header(headers, b"whatever", ["different thing"])
 
     assert headers == [
         (b"connection", b"close"),

h11/tests/test_io.py

@@ -1,7 +1,7 @@
 import pytest
 
 from .._events import *
-from .._headers import normalize_and_validate
+from .._headers import normalize_and_validate, Headers
 from .._readers import (
     _obsolete_line_fold,
     ChunkedReader,
@@ -121,7 +121,7 @@ def test_writers_unusual():
         normalize_and_validate([("foo", "bar"), ("baz", "quux")]),
         b"foo: bar\r\nbaz: quux\r\n\r\n",
     )
-    tw(write_headers, [], b"\r\n")
+    tw(write_headers, Headers([]), b"\r\n")
 
     # We understand HTTP/1.0, but we don't speak it
     with pytest.raises(LocalProtocolError):