Support chunked responses with normal reads.

Author: Lukasa

hyper/http11/response.py

@@ -74,6 +74,9 @@ def __init__(self, code, reason, headers, sock):
         else:
             self._decompressobj = None
 
+        self._buffered_data = b''
+        self._chunker = None
+
     def read(self, amt=None, decode_content=True):
         """
         Reads the response body, or up to the next ``amt`` bytes.
@@ -90,8 +93,32 @@ def read(self, amt=None, decode_content=True):
         if self._sock is None:
             return b''
 
-        # FIXME: Handle chunked transfer encoding
-        assert not self._chunked
+        if self._chunked:
+            # If we're doing a full read, read it as chunked and then just join
+            # the chunks together!
+            if amt is None:
+                return self._buffered_data + b''.join(self.read_chunked())
+
+            if self._chunker is None:
+                self._chunker = self.read_chunked()
+
+            # Otherwise, we have a certain amount of data we want to read.
+            current_amount = len(self._buffered_data)
+
+            extra_data = [self._buffered_data]
+            while current_amount < amt:
+                try:
+                    chunk = next(self._chunker)
+                except StopIteration:
+                    self.close()
+                    break
+
+                current_amount += len(chunk)
+                extra_data.append(chunk)
+
+            data = b''.join(extra_data)
+            self._buffered_data = data[amt:]
+            return data[:amt]
 
         # If we're asked to do a read without a length, we need to read
         # everything. That means either the entire content length, or until the

test/test_http11.py

@@ -445,6 +445,43 @@ def test_response_transparently_decrypts_chunked_gzip(self):
 
         assert received_body == b'this is test data'
 
+    def test_chunked_normal_read(self):
+        d = DummySocket()
+        r = HTTP11Response(200, 'OK', {b'transfer-encoding': [b'chunked']}, d)
+
+        data = (
+            b'4\r\nwell\r\n'
+            b'4\r\nwell\r\n'
+            b'4\r\nwhat\r\n'
+            b'4\r\nhave\r\n'
+            b'2\r\nwe\r\n'
+            b'a\r\nhereabouts\r\n'
+            b'0\r\n\r\n'
+        )
+        d._buffer = BytesIO(data)
+
+        assert r.read() == b'wellwellwhathavewehereabouts'
+
+    def test_chunk_length_read(self):
+        d = DummySocket()
+        r = HTTP11Response(200, 'OK', {b'transfer-encoding': [b'chunked']}, d)
+
+        data = (
+            b'4\r\nwell\r\n'
+            b'4\r\nwell\r\n'
+            b'4\r\nwhat\r\n'
+            b'4\r\nhave\r\n'
+            b'2\r\nwe\r\n'
+            b'a\r\nhereabouts\r\n'
+            b'0\r\n\r\n'
+        )
+        d._buffer = BytesIO(data)
+
+        assert r.read(5) == b'wellw'
+        assert r.read(15) == b'ellwhathavewehe'
+        assert r.read(20) == b'reabouts'
+        assert r.read() == b''
+
 class DummySocket(object):
     def __init__(self):
         self.queue = []