Skip to content

Commit

Permalink
fix: add correct support for compressing file-like objects
Browse files Browse the repository at this point in the history
Signed-off-by: Norbert Biczo <[email protected]>
  • Loading branch information
pyrooka committed Sep 25, 2023
1 parent 17039a3 commit 35216f3
Show file tree
Hide file tree
Showing 2 changed files with 82 additions and 3 deletions.
10 changes: 7 additions & 3 deletions ibm_cloud_sdk_core/base_service.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
# limitations under the License.

import gzip
import io
import json as json_import
import logging
import platform
Expand All @@ -39,6 +40,7 @@
read_external_sources,
strip_extra_slashes,
SSLHTTPAdapter,
GzipStream,
)
from .version import __version__

Expand Down Expand Up @@ -420,10 +422,12 @@ def prepare_request(
# Compress the request body if applicable
if self.get_enable_gzip_compression() and 'content-encoding' not in headers and request['data'] is not None:
headers['content-encoding'] = 'gzip'
uncompressed_data = request['data']
request_body = gzip.compress(uncompressed_data)
request['data'] = request_body
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)

# Next, we need to process the 'files' argument to try to fill in
# any missing filenames where possible.
Expand Down
75 changes: 75 additions & 0 deletions ibm_cloud_sdk_core/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,8 @@
# limitations under the License.
# from ibm_cloud_sdk_core.authenticators import Authenticator
import datetime
import gzip
import io
import json as json_import
import re
import ssl
Expand Down Expand Up @@ -43,6 +45,79 @@ def init_poolmanager(self, connections, maxsize, block):
super().init_poolmanager(connections, maxsize, block, ssl_context=ssl_context)


class GzipStream(io.IOBase):
"""Compress files on the fly.
GzipStream is a helper class around the gzip library. It helps to
compress already opened files (file-like objects) on the fly, so
there is no need to read everything into the memory and call the
`compress` function on it.
The GzipFile is opened on the instance itself so it needs to act
as a file-like object.
Args:
input: a file-like object to be compressed
"""

def __init__(self, source: io.IOBase):
self.uncompressed = source
self.buffer = b''

self.compressor = gzip.GzipFile(fileobj=self, mode='wb')

def read(self, size: int = -1):
"""Compresses and returns the reqested size of data.
Args:
size: how many bytes to return. -1 to read and compress the whole file
"""
if (size < 0) or (len(self.buffer) < size):
for raw in self.uncompressed:
# We need to encode text like streams (e.g. TextIOWrapper) to bytes.
if isinstance(raw, str):
raw = raw.encode()

self.compressor.write(raw)

# Stop compressing if we reached the max allowed size.
if 0 < size < len(self.buffer):
self.compressor.flush()
break
else:
self.compressor.close()

if size < 0:
# Return all data from the buffer.
compressed = self.buffer
self.buffer = b''
else:
# If we already have enough data in our buffer
# return the desired chunk of bytes
compressed = self.buffer[:size]
# then remove them from the buffer.
self.buffer = self.buffer[size:]

return compressed

def flush(self):
"""Not implemented."""
# Since this "pipe" sits between 2 other stream (source/read -> target/write)
# it wouldn't be worth to implemet flushing.
pass

def write(self, compressed: bytes):
"""Appened the compressed data to the buffer
This happens when the target stream calls the `read` method and
that triggers the gzip "compressor".
"""
self.buffer += compressed

def close(self):
"""Closes the underlying file-like object."""
self.uncompressed.close()


def has_bad_first_or_last_char(val: str) -> bool:
"""Returns true if a string starts with any of: {," ; or ends with any of: },".
Expand Down

0 comments on commit 35216f3

Please sign in to comment.