Don't leak compressed stream (#630)

CHANGELOG.md

@@ -3,6 +3,7 @@
 - Work around changes to `urllib.parse.urlsplit` (PR [#633](https://github.com/RaRe-Technologies/smart_open/pull/633), [@judahrand](https://github.com/judahrand)
 - Change python_requires version to fix PEP 440 issue (PR [#639](https://github.com/RaRe-Technologies/smart_open/pull/639), [@lucasvieirasilva](https://github.com/lucasvieirasilva))
 - New blob_properties transport parameter for GCS (PR [#632](https://github.com/RaRe-Technologies/smart_open/pull/632), [@FHTheron](https://github.com/FHTheron))
+- Don't leak compressed stream (PR [#636](https://github.com/RaRe-Technologies/smart_open/pull/636), [@ampanasiuk](https://github.com/ampanasiuk))
 
 # 5.1.0, 25 May 2021
 

smart_open/compression.py

@@ -67,14 +67,40 @@ def register_compressor(ext, callback):
     _COMPRESSOR_REGISTRY[ext] = callback
 
 
+def tweak_close(outer, inner):
+    """Ensure that closing the `outer` stream closes the `inner` stream as well.
+
+    Use this when your compression library's `close` method does not
+    automatically close the underlying filestream.  See
+    https://github.com/RaRe-Technologies/smart_open/issues/630 for an
+    explanation why that is a problem for smart_open.
+    """
+    outer_close = outer.close
+
+    def close_both(*args):
+        nonlocal inner
+        try:
+            outer_close()
+        finally:
+            if inner:
+                inner, fp = None, inner
+                fp.close()
+
+    outer.close = close_both
+
+
 def _handle_bz2(file_obj, mode):
     from bz2 import BZ2File
-    return BZ2File(file_obj, mode)
+    result = BZ2File(file_obj, mode)
+    tweak_close(result, file_obj)
+    return result
 
 
 def _handle_gzip(file_obj, mode):
     import gzip
-    return gzip.GzipFile(fileobj=file_obj, mode=mode)
+    result = gzip.GzipFile(fileobj=file_obj, mode=mode)
+    tweak_close(result, file_obj)
+    return result
 
 
 def compression_wrapper(file_obj, mode, compression):

smart_open/tests/test_smart_open.py

@@ -453,7 +453,7 @@ def _test_compressed_http(self, suffix, query):
         buf = make_buffer(name='data' + suffix)
         with smart_open.open(buf, 'wb') as outfile:
             outfile.write(raw_data)
-        compressed_data = buf.getvalue()
+        compressed_data = buf._value_when_closed
         # check that the string was actually compressed
         self.assertNotEqual(compressed_data, raw_data)
 
@@ -494,8 +494,17 @@ def make_buffer(cls=io.BytesIO, initial_value=None, name=None, noclose=False):
     buf = cls(initial_value) if initial_value else cls()
     if name is not None:
         buf.name = name
-    if noclose:
-        buf.close = lambda: None
+
+    buf._value_when_closed = None
+    orig_close = buf.close
+
+    def close():
+        if buf.close.call_count == 1:
+            buf._value_when_closed = buf.getvalue()
+        if not noclose:
+            orig_close()
+
+    buf.close = mock.Mock(side_effect=close)
     return buf
 
 
@@ -652,7 +661,7 @@ def test_name_write(self):
         buf = make_buffer(name='data.bz2')
         with smart_open.open(buf, 'wb') as sf:
             sf.write(data)
-        self.assertEqual(bz2.decompress(buf.getvalue()), data)
+        self.assertEqual(bz2.decompress(buf._value_when_closed), data)
 
     def test_open_side_effect(self):
         """
@@ -1496,16 +1505,32 @@ def write_callback(request):
         assert responses.calls[3].request.url == "http://127.0.0.1:8440/file"
 
 
-class CompressionFormatTest(unittest.TestCase):
+_DECOMPRESSED_DATA = "не слышны в саду даже шорохи".encode("utf-8")
+_MOCK_TIME = mock.Mock(return_value=1620256567)
+
+
+def gzip_compress(data, filename=None):
+    #
+    # gzip.compress is sensitive to the current time and the destination filename.
+    # This function fixes those variables for consistent compression results.
+    #
+    buf = io.BytesIO()
+    buf.name = filename
+    with mock.patch('time.time', _MOCK_TIME):
+        gzip.GzipFile(fileobj=buf, mode='w').write(data)
+    return buf.getvalue()
+
+
+class CompressionFormatTest(parameterizedtestcase.ParameterizedTestCase):
     """Test transparent (de)compression."""
 
     def write_read_assertion(self, suffix):
         test_file = make_buffer(name='file' + suffix)
         with smart_open.open(test_file, 'wb') as fout:
             fout.write(SAMPLE_BYTES)
-        self.assertNotEqual(SAMPLE_BYTES, test_file.getvalue())
+        self.assertNotEqual(SAMPLE_BYTES, test_file._value_when_closed)
         # we have to recreate the buffer because it is closed
-        test_file = make_buffer(initial_value=test_file.getvalue(), name=test_file.name)
+        test_file = make_buffer(initial_value=test_file._value_when_closed, name=test_file.name)
         with smart_open.open(test_file, 'rb') as fin:
             self.assertEqual(fin.read(), SAMPLE_BYTES)
 
@@ -1517,6 +1542,20 @@ def test_open_gz(self):
         m = hashlib.md5(data)
         assert m.hexdigest() == '18473e60f8c7c98d29d65bf805736a0d', 'Failed to read gzip'
 
+    @parameterizedtestcase.ParameterizedTestCase.parameterize(
+        ("extension", "compressed"),
+        [
+            (".gz", gzip_compress(_DECOMPRESSED_DATA, 'key')),
+            (".bz2", bz2.compress(_DECOMPRESSED_DATA)),
+        ],
+    )
+    def test_closes_compressed_stream(self, extension, compressed):
+        """Transparent compression closes the compressed stream?"""
+        compressed_stream = make_buffer(initial_value=compressed, name=f"file{extension}")
+        with smart_open.open(compressed_stream, encoding="utf-8"):
+            pass
+        assert compressed_stream.close.call_count == 1
+
     def test_write_read_gz(self):
         """Can write and read gzip?"""
         self.write_read_assertion('.gz')
@@ -1841,22 +1880,6 @@ def test(self):
         self.assertEqual(expected, actual)
 
 
-_DECOMPRESSED_DATA = "не слышны в саду даже шорохи".encode("utf-8")
-_MOCK_TIME = mock.Mock(return_value=1620256567)
-
-
-def gzip_compress(data, filename=None):
-    #
-    # gzip.compress is sensitive to the current time and the destination filename.
-    # This function fixes those variables for consistent compression results.
-    #
-    buf = io.BytesIO()
-    buf.name = filename
-    with mock.patch('time.time', _MOCK_TIME):
-        gzip.GzipFile(fileobj=buf, mode='w').write(data)
-    return buf.getvalue()
-
-
 @mock_s3
 @mock.patch('time.time', _MOCK_TIME)
 class S3CompressionTestCase(parameterizedtestcase.ParameterizedTestCase):