Support .flush for Compressor (#68)

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "cramjam"
-version = "2.4.0-rc1"
+version = "2.4.0-rc2"
 authors = ["Miles Granger <miles59923@gmail.com>"]
 edition = "2018"
 license = "MIT License"

src/brotli.rs

@@ -77,6 +77,11 @@ impl Compressor {
         crate::io::stream_compress(&mut self.inner, input)
     }
 
+    /// Flush and return current compressed stream
+    pub fn flush(&mut self) -> PyResult<RustyBuffer> {
+        crate::io::stream_flush(&mut self.inner, |e| e.get_mut())
+    }
+
     /// Consume the current compressor state and return the compressed stream
     /// **NB** The compressor will not be usable after this method is called.
     pub fn finish(&mut self) -> PyResult<RustyBuffer> {

src/deflate.rs

@@ -78,6 +78,11 @@ impl Compressor {
         crate::io::stream_compress(&mut self.inner, input)
     }
 
+    /// Flush and return current compressed stream
+    pub fn flush(&mut self) -> PyResult<RustyBuffer> {
+        crate::io::stream_flush(&mut self.inner, |e| e.get_mut())
+    }
+
     /// Consume the current compressor state and return the compressed stream
     /// **NB** The compressor will not be usable after this method is called.
     pub fn finish(&mut self) -> PyResult<RustyBuffer> {

src/gzip.rs

@@ -77,6 +77,11 @@ impl Compressor {
         crate::io::stream_compress(&mut self.inner, input)
     }
 
+    /// Flush and return current compressed stream
+    pub fn flush(&mut self) -> PyResult<RustyBuffer> {
+        crate::io::stream_flush(&mut self.inner, |e| e.get_mut())
+    }
+
     /// Consume the current compressor state and return the compressed stream
     /// **NB** The compressor will not be usable after this method is called.
     pub fn finish(&mut self) -> PyResult<RustyBuffer> {

src/io.rs

@@ -635,3 +635,22 @@ where
         None => Ok(RustyBuffer::from(vec![])),
     }
 }
+
+// flush inner encoder data out
+pub(crate) fn stream_flush<W, F>(encoder: &mut Option<W>, cursor_mut_ref: F) -> PyResult<RustyBuffer>
+where
+    W: Write,
+    F: Fn(&mut W) -> &mut Cursor<Vec<u8>>,
+{
+    match encoder {
+        Some(inner) => {
+            crate::to_py_err!(CompressionError -> inner.flush())?;
+            let cursor = cursor_mut_ref(inner);
+            let buf = RustyBuffer::from(cursor.get_ref().clone());
+            cursor.get_mut().truncate(0);
+            cursor.set_position(0);
+            Ok(buf)
+        }
+        None => Ok(RustyBuffer::from(vec![])),
+    }
+}

src/lz4.rs

@@ -143,6 +143,15 @@ impl Compressor {
         crate::io::stream_compress(&mut self.inner, input)
     }
 
+    /// Flush and return current compressed stream
+    #[allow(mutable_transmutes)] // TODO: feature req to lz4 to get mut ref to writer
+    pub fn flush(&mut self) -> PyResult<RustyBuffer> {
+        crate::io::stream_flush(&mut self.inner, |e| {
+            let writer = e.writer();
+            unsafe { std::mem::transmute::<&Cursor<Vec<u8>>, &mut Cursor<Vec<u8>>>(writer) }
+        })
+    }
+
     /// Consume the current compressor state and return the compressed stream
     /// **NB** The compressor will not be usable after this method is called.
     pub fn finish(&mut self) -> PyResult<RustyBuffer> {

src/snappy.rs

@@ -144,6 +144,11 @@ impl Compressor {
         crate::io::stream_compress(&mut self.inner, input)
     }
 
+    /// Flush and return current compressed stream
+    pub fn flush(&mut self) -> PyResult<RustyBuffer> {
+        crate::io::stream_flush(&mut self.inner, |e| e.get_mut())
+    }
+
     /// Consume the current compressor state and return the compressed stream
     /// **NB** The compressor will not be usable after this method is called.
     pub fn finish(&mut self) -> PyResult<RustyBuffer> {

src/zstd.rs

@@ -76,6 +76,11 @@ impl Compressor {
         crate::io::stream_compress(&mut self.inner, input)
     }
 
+    /// Flush and return current compressed stream
+    pub fn flush(&mut self) -> PyResult<RustyBuffer> {
+        crate::io::stream_flush(&mut self.inner, |e| e.get_mut())
+    }
+
     /// Consume the current compressor state and return the compressed stream
     /// **NB** The compressor will not be usable after this method is called.
     pub fn finish(&mut self) -> PyResult<RustyBuffer> {

tests/test_variants.py

@@ -254,9 +254,14 @@ def test_gzip_multiple_streams():
 )
 def test_streams_compressor(mod):
     compressor = mod.Compressor()
+
     compressor.compress(b"foo")
+    out = bytes(compressor.flush())
+
     compressor.compress(b"bar")
-    out = compressor.finish()
+    out += bytes(compressor.flush())
+
+    out += bytes(compressor.finish())
     decompressed = mod.decompress(out)
     assert bytes(decompressed) == b"foobar"
 
@@ -266,4 +271,4 @@ def test_streams_compressor(mod):
 
     # compress will raise an error as the stream is completed
     with pytest.raises(cramjam.CompressionError):
-        compressor.compress(b'data')
+        compressor.compress(b"data")