From bdc27d3d2a3105f9b524e4539b3845d221f22cdc Mon Sep 17 00:00:00 2001 From: Norbert Biczo Date: Tue, 26 Sep 2023 13:32:15 +0200 Subject: [PATCH] chore: move all compression logic to helper class Signed-off-by: Norbert Biczo --- ibm_cloud_sdk_core/base_service.py | 8 +------- ibm_cloud_sdk_core/utils.py | 16 +++++++++++++--- test/test_base_service.py | 10 +++++----- 3 files changed, 19 insertions(+), 15 deletions(-) diff --git a/ibm_cloud_sdk_core/base_service.py b/ibm_cloud_sdk_core/base_service.py index 0003fa9..06aef87 100644 --- a/ibm_cloud_sdk_core/base_service.py +++ b/ibm_cloud_sdk_core/base_service.py @@ -14,8 +14,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import gzip -import io import json as json_import import logging import platform @@ -423,11 +421,7 @@ def prepare_request( if self.get_enable_gzip_compression() and 'content-encoding' not in headers and request['data'] is not None: headers['content-encoding'] = 'gzip' request['headers'] = headers - raw = request['data'] - # Handle the compression for file-like objects. - # We need to use a custom stream/pipe method to prevent - # reading the whole file into the memory. - request['data'] = GzipStream(raw) if isinstance(raw, io.IOBase) else gzip.compress(raw) + request['data'] = GzipStream(request['data']) # Next, we need to process the 'files' argument to try to fill in # any missing filenames where possible. diff --git a/ibm_cloud_sdk_core/utils.py b/ibm_cloud_sdk_core/utils.py index f3eb43c..e961b05 100644 --- a/ibm_cloud_sdk_core/utils.py +++ b/ibm_cloud_sdk_core/utils.py @@ -56,13 +56,23 @@ class GzipStream(io.IOBase): as a file-like object. Args: - input: a file-like object to be compressed + input: the source of the data to be compressed. + It can be a file-like object, bytes or string. """ - def __init__(self, source: io.IOBase): - self.uncompressed = source + def __init__(self, source: Union[io.IOBase, bytes, str]): self.buffer = b'' + if isinstance(source, io.IOBase): + # The input is already a file-like object, use it as-is. + self.uncompressed = source + elif isinstance(source, str): + # Strings must be handled with StringIO. + self.uncompressed = io.StringIO(source) + else: + # Handle the rest as raw bytes. + self.uncompressed = io.BytesIO(source) + self.compressor = gzip.GzipFile(fileobj=self, mode='wb') def read(self, size: int = -1): diff --git a/test/test_base_service.py b/test/test_base_service.py index a3527f5..349f083 100644 --- a/test/test_base_service.py +++ b/test/test_base_service.py @@ -607,13 +607,13 @@ def test_gzip_compression(): service.set_enable_gzip_compression(True) assert service.get_enable_gzip_compression() prepped = service.prepare_request('GET', url='', data=json.dumps({"foo": "bar"})) - assert prepped['data'] == gzip.compress(b'{"foo": "bar"}') + assert prepped['data'].read() == gzip.compress(b'{"foo": "bar"}') assert prepped['headers'].get('content-encoding') == 'gzip' # Should return compressed data when gzip is on for non-json data assert service.get_enable_gzip_compression() prepped = service.prepare_request('GET', url='', data=b'rawdata') - assert prepped['data'] == gzip.compress(b'rawdata') + assert prepped['data'].read() == gzip.compress(b'rawdata') assert prepped['headers'].get('content-encoding') == 'gzip' # Should return compressed data when gzip is on for gzip file data @@ -624,7 +624,7 @@ def test_gzip_compression(): with gzip.GzipFile(mode='rb', fileobj=t_f) as gz_f: gzip_data = gz_f.read() prepped = service.prepare_request('GET', url='', data=gzip_data) - assert prepped['data'] == gzip.compress(t_f.read()) + assert prepped['data'].read() == gzip.compress(t_f.read()) assert prepped['headers'].get('content-encoding') == 'gzip' # Should return compressed json data when gzip is on for gzip file json data @@ -635,7 +635,7 @@ def test_gzip_compression(): with gzip.GzipFile(mode='rb', fileobj=t_f) as gz_f: gzip_data = gz_f.read() prepped = service.prepare_request('GET', url='', data=gzip_data) - assert prepped['data'] == gzip.compress(t_f.read()) + assert prepped['data'].read() == gzip.compress(t_f.read()) assert prepped['headers'].get('content-encoding') == 'gzip' # Should return uncompressed data when content-encoding is set @@ -683,7 +683,7 @@ def test_gzip_compression_external(): assert service.service_url == 'https://mockurl' assert service.get_enable_gzip_compression() is True prepped = service.prepare_request('GET', url='', data=json.dumps({"foo": "bar"})) - assert prepped['data'] == gzip.compress(b'{"foo": "bar"}') + assert prepped['data'].read() == gzip.compress(b'{"foo": "bar"}') assert prepped['headers'].get('content-encoding') == 'gzip'