Turn read_chunked into a generator.

While we're there, add support for transparent content
decompression.

Author: Lukasa

hyper/http11/response.py

@@ -154,10 +154,15 @@ def read(self, amt=None, decode_content=True):
 
         return data
 
-    def read_chunked(self):
+    def read_chunked(self, decode_content=True):
         """
-        Reads chunked transfer encoded bodies. Each read returns a single
-        chunk.
+        Reads chunked transfer encoded bodies. This method returns a generator:
+        each iteration of which yields one chunk *unless* the chunks are
+        compressed, in which case it yields whatever the decompressor provides
+        for each chunk.
+
+        .. warning:: This may yield the empty string, without that being the
+                     end of the body!
         """
         if not self._chunked:
             raise ChunkedDecodeError(
@@ -166,31 +171,42 @@ def read_chunked(self):
 
         # Return early if possible.
         if self._sock is None:
-            return b''
+            return
 
-        # Read to the newline to get the chunk length. This is a hexadecimal
-        # integer.
-        chunk_length = int(self._sock.readline().tobytes().strip(), 16)
+        while True:
+            # Read to the newline to get the chunk length. This is a
+            # hexadecimal integer.
+            chunk_length = int(self._sock.readline().tobytes().strip(), 16)
+            data = b''
 
-        # If the chunk length is zero, consume the newline and then we're done.
-        if not chunk_length:
-            self._sock.readline()
-            return b''
+            # If the chunk length is zero, consume the newline and then we're
+            # done. If we were decompressing data, return the remaining data.
+            if not chunk_length:
+                self._sock.readline()
 
-        # Then read that many bytes.
-        data = b''
+                if decode_content and self._decompressobj:
+                    yield self._decompressobj.flush()
 
-        while chunk_length > 0:
-            chunk = self._sock.recv(chunk_length).tobytes()
-            data += chunk
-            chunk_length -= len(chunk)
+                break
 
-        assert chunk_length == 0
+            # Then read that many bytes.
+            while chunk_length > 0:
+                chunk = self._sock.recv(chunk_length).tobytes()
+                data += chunk
+                chunk_length -= len(chunk)
 
-        # Now, consume the newline.
-        self._sock.readline()
+            assert chunk_length == 0
 
-        return data
+            # Now, consume the newline.
+            self._sock.readline()
+
+            # We may need to decode the body.
+            if decode_content and self._decompressobj and data:
+                data = self._decompressobj.decompress(data)
+
+            yield data
+
+        return
 
     def close(self):
         """

test/test_http11.py

@@ -402,24 +402,48 @@ def test_basic_chunked_read(self):
             b'well', b'well', b'what', b'have', b'we', b'hereabouts'
         ]
 
-        for c in results:
-            assert r.read_chunked() == c
+        for c1, c2 in zip(results, r.read_chunked()):
+            assert c1 == c2
 
-        assert not r.read_chunked()
+        assert not list(r.read_chunked())
 
     def test_chunked_read_of_non_chunked(self):
         r = HTTP11Response(200, 'OK', {b'content-length': [b'0']}, None)
 
         with pytest.raises(ChunkedDecodeError):
-            r.read_chunked()
+            list(r.read_chunked())
 
     def test_chunked_read_aborts_early(self):
         r = HTTP11Response(
             200, 'OK', {b'transfer-encoding': [b'chunked']}, None
         )
 
-        assert not r.read_chunked()
+        assert not list(r.read_chunked())
 
+    def test_response_transparently_decrypts_chunked_gzip(self):
+        d = DummySocket()
+        headers = {
+            b'content-encoding': [b'gzip'],
+            b'transfer-encoding': [b'chunked'],
+        }
+        r = HTTP11Response(200, 'OK', headers, d)
+
+        c = zlib_compressobj(wbits=24)
+        body = c.compress(b'this is test data')
+        body += c.flush()
+
+        data = b''
+        for index in range(0, len(body), 2):
+            data += '2\r\n' + body[index:index+2] + '\r\n'
+
+        data += b'0\r\n\r\n'
+        d._buffer = BytesIO(data)
+
+        received_body = b''
+        for chunk in r.read_chunked():
+            received_body += chunk
+
+        assert received_body == b'this is test data'
 
 class DummySocket(object):
     def __init__(self):