Support chunked reads of HTTP/2.

Author: Lukasa

hyper/http20/response.py

@@ -135,6 +135,34 @@ def read(self, amt=None, decode_content=True):
 
         return data
 
+    def read_chunked(self, decode_content=True):
+        """
+        Reads chunked transfer encoded bodies. This method returns a generator:
+        each iteration of which yields one data frame *unless* the frames
+        contain compressed data and ``decode_content`` is ``True``, 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!
+        """
+        while True:
+            data = self._stream._read_one_frame()
+
+            if data is None:
+                break
+
+            if decode_content and self._decompressobj:
+                data = self._decompressobj.decompress(data)
+
+            yield data
+
+        if decode_content and self._decompressobj:
+            yield self._decompressobj.flush()
+
+        self.close()
+
+        return
+
     def fileno(self):
         """
         Return the ``fileno`` of the underlying socket. This function is

hyper/http20/stream.py

@@ -179,6 +179,19 @@ def listlen(list):
         self.data = []
         return result
 
+    def _read_one_frame(self):
+        """
+        Reads a single data frame from the stream and returns it.
+        """
+        # Keep reading until the stream is closed or we have a data frame.
+        while not self._remote_closed and not self.data:
+            self._recv_cb()
+
+        try:
+            return self.data.pop(0)
+        except IndexError:
+            return None
+
     def receive_frame(self, frame):
         """
         Handle a frame received on this stream.

test/test_hyper.py

@@ -1861,6 +1861,44 @@ def inner():
         assert s.gettrailers() == HTTPHeaderMap(trailers)
         assert s.data == [b'testdata']
 
+    def test_can_read_single_frames_from_streams(self):
+        out_frames = []
+        in_frames = []
+
+        def send_cb(frame, tolerate_peer_gone=False):
+            out_frames.append(frame)
+
+        def recv_cb(s):
+            def inner():
+                s.receive_frame(in_frames.pop(0))
+            return inner
+
+        s = Stream(1, send_cb, None, None, None, None, FlowControlManager(800))
+        s._recv_cb = recv_cb(s)
+        s.state = STATE_HALF_CLOSED_LOCAL
+
+        # Provide two data frames to read.
+        f = DataFrame(1)
+        f.data = b'hi there!'
+        in_frames.append(f)
+
+        f = DataFrame(1)
+        f.data = b'hi there again!'
+        f.flags.add('END_STREAM')
+        in_frames.append(f)
+
+        data = s._read_one_frame()
+        assert data == b'hi there!'
+
+        data = s._read_one_frame()
+        assert data == b'hi there again!'
+
+        data = s._read_one_frame()
+        assert data is None
+
+        data = s._read()
+        assert data == b''
+
 
 class TestResponse(object):
     def test_status_is_stripped_from_headers(self):
@@ -1968,6 +2006,33 @@ def test_trailers_are_read(self):
         assert resp.trailers['a'] == [b'b']
         assert resp.trailers['c'] == [b'd']
 
+    def test_read_frames(self):
+        headers = HTTPHeaderMap([(':status', '200')])
+        stream = DummyStream(None)
+        chunks = [b'12', b'3456', b'78', b'9']
+        stream.data_frames = chunks
+        resp = HTTP20Response(headers, stream)
+
+        for recv, expected in zip(resp.read_chunked(), chunks[:]):
+            assert recv == expected
+
+    def test_read_compressed_frames(self):
+        headers = HTTPHeaderMap([(':status', '200'), ('content-encoding', 'gzip')])
+        c = zlib_compressobj(wbits=24)
+        body = c.compress(b'this is test data')
+        body += c.flush()
+
+        stream = DummyStream(None)
+        chunks = [body[x:x+2] for x in range(0, len(body), 2)]
+        stream.data_frames = chunks
+        resp = HTTP20Response(headers, stream)
+
+        received = b''
+        for chunk in resp.read_chunked():
+            received += chunk
+
+        assert received == b'this is test data'
+
 
 class TestHTTP20Adapter(object):
     def test_adapter_reuses_connections(self):
@@ -2090,6 +2155,7 @@ def recv(self, l):
 class DummyStream(object):
     def __init__(self, data, trailers=None):
         self.data = data
+        self.data_frames = []
         self.closed = False
         self.response_headers = {}
         self._remote_closed = False
@@ -2112,6 +2178,12 @@ def _read(self, *args, **kwargs):
 
         return d
 
+    def _read_one_frame(self):
+        try:
+            return self.data_frames.pop(0)
+        except IndexError:
+            return None
+
     def close(self):
         if not self.closed:
             self.closed = True