diff --git a/src/caosdb/connection/connection.py b/src/caosdb/connection/connection.py index a81c9a5b890069eea9cf01a180a9a1ccc84aeb2b..aaa09dfe9d804a5465f322bfe8db5826b304211b 100644 --- a/src/caosdb/connection/connection.py +++ b/src/caosdb/connection/connection.py @@ -65,6 +65,9 @@ class _WrappedHTTPResponse(CaosDBHTTPResponse): def __init__(self, response): self.response = response + self._generator = None + self._buffer = b'' + self._stream_consumed = False @property def reason(self): @@ -75,12 +78,44 @@ class _WrappedHTTPResponse(CaosDBHTTPResponse): return self.response.status_code def read(self, size=None): - if size is None or size == 0: - result = b'' - for chunk in self.response.iter_content(chunk_size=1024): - result = result + chunk + if self._stream_consumed is True: + raise RuntimeError("Stream is consumed") + + if self._buffer is None: + # the buffer has been drained in the previous call. + self._stream_consumed = True + return b'' + + if self._generator is None and (size is None or size == 0): + # return full content at once + self._stream_consumed = True + return self.response.content + + if len(self._buffer) >= size: + # still enough bytes in the buffer + result = chunk[:size] + self._buffer = chunk[size:] + return result + + if self._generator is None: + # first call to this method + if size is None or size == 0: + size = 512 + self._generator = self.response.iter_content(size) + + try: + # read new data into the buffer + chunk = self._buffer + next(self._generator) + result = chunk[:size] + if len(result) == 0: + self._stream_consumed = True + self._buffer = chunk[size:] + return result + except StopIteration: + # drain buffer + result = self._buffer + self._buffer = None return result - return self.response.iter_content(size) def getheader(self, name, default=None): return self.response.headers[name] if name in self.response.headers else default