From 837dccbdd14b900a0e5056a0410606f9b7dde9ac Mon Sep 17 00:00:00 2001
From: Timm Fitschen <t.fitschen@indiscale.com>
Date: Wed, 9 Nov 2022 12:18:26 +0100
Subject: [PATCH] FIX: implementation of WrappedResponse
---
src/caosdb/connection/connection.py | 45 +++++++++++++++++++++++++----
1 file changed, 40 insertions(+), 5 deletions(-)
diff --git a/src/caosdb/connection/connection.py b/src/caosdb/connection/connection.py
index a81c9a5b..aaa09dfe 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
--
GitLab