Skip to content

Commit

Permalink
Not able to upload a compressed S3 file using close() (#838)
Browse files Browse the repository at this point in the history
* Call inner close() call during outer close() call in FileLikeProxy class

* Return close() result in FileLikeProxy class
  • Loading branch information
jbarragan-bridge authored Oct 4, 2024
1 parent 3c610c3 commit fed7b78
Show file tree
Hide file tree
Showing 2 changed files with 44 additions and 0 deletions.
37 changes: 37 additions & 0 deletions smart_open/tests/test_s3.py
Original file line number Diff line number Diff line change
Expand Up @@ -634,6 +634,43 @@ def test_writebuffer(self):

assert actual == contents

def test_write_gz_using_context_manager(self):
"""Does s3 multipart upload create a compressed file using context manager?"""
contents = b'get ready for a surprise'
with smart_open.open(
f's3://{BUCKET_NAME}/{WRITE_KEY_NAME}.gz',
mode="wb",
transport_params={
"multipart_upload": True,
"min_part_size": 10,
}
) as fout:
fout.write(contents)

with smart_open.open(f's3://{BUCKET_NAME}/{WRITE_KEY_NAME}.gz', 'rb') as fin:
actual = fin.read()

assert actual == contents

def test_write_gz_not_using_context_manager(self):
"""Does s3 multipart upload create a compressed file not using context manager but close()?"""
contents = b'get ready for a surprise'
fout = smart_open.open(
f's3://{BUCKET_NAME}/{WRITE_KEY_NAME}.gz',
mode="wb",
transport_params={
"multipart_upload": True,
"min_part_size": 10,
}
)
fout.write(contents)
fout.close()

with smart_open.open(f's3://{BUCKET_NAME}/{WRITE_KEY_NAME}.gz', 'rb') as fin:
actual = fin.read()

assert actual == contents

def test_write_gz_with_error(self):
"""Does s3 multipart upload abort for a failed compressed file upload?"""
with self.assertRaises(ValueError):
Expand Down
7 changes: 7 additions & 0 deletions smart_open/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -223,3 +223,10 @@ def __exit__(self, *args, **kwargs):

def __next__(self):
return self.__wrapped__.__next__()

def close(self):
try:
return self.__wrapped__.close()
finally:
if self.__inner != self.__wrapped__: # Don't close again if inner and wrapped are the same
self.__inner.close()

0 comments on commit fed7b78

Please sign in to comment.