Fixup: refactor finish into io module

src/brotli.rs

@@ -80,16 +80,7 @@ impl Compressor {
     /// 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> {
-        let mut inner = None;
-        std::mem::swap(&mut inner, &mut self.inner);
-
-        match inner {
-            Some(inner) => {
-                let result = to_py_err!(CompressionError -> inner.finish())?;
-                Ok(RustyBuffer::from(result.into_inner()))
-            }
-            None => Ok(RustyBuffer::from(vec![])),
-        }
+        crate::io::stream_finish(&mut self.inner, |inner| inner.finish().map(|c| c.into_inner()))
     }
 }
 

src/deflate.rs

@@ -81,16 +81,7 @@ impl Compressor {
     /// 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> {
-        let mut inner = None;
-        std::mem::swap(&mut inner, &mut self.inner);
-
-        match inner {
-            Some(inner) => {
-                let result = to_py_err!(CompressionError -> inner.finish())?;
-                Ok(RustyBuffer::from(result.into_inner()))
-            }
-            None => Ok(RustyBuffer::from(vec![])),
-        }
+        crate::io::stream_finish(&mut self.inner, |inner| inner.finish().map(|c| c.into_inner()))
     }
 }
 

src/gzip.rs

@@ -80,16 +80,7 @@ impl Compressor {
     /// 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> {
-        let mut inner = None;
-        std::mem::swap(&mut inner, &mut self.inner);
-
-        match inner {
-            Some(inner) => {
-                let result = to_py_err!(CompressionError -> inner.finish())?;
-                Ok(RustyBuffer::from(result.into_inner()))
-            }
-            None => Ok(RustyBuffer::from(vec![])),
-        }
+        crate::io::stream_finish(&mut self.inner, |inner| inner.finish().map(|c| c.into_inner()))
     }
 }
 

src/io.rs

@@ -600,10 +600,10 @@ impl Read for RustyFile {
 // general stream compression interface. Can't use associated types due to pyo3::pyclass
 // not supporting generic structs.
 #[inline(always)]
-pub(crate) fn stream_compress<W: Write>(inner: &mut Option<W>, input: &[u8]) -> PyResult<usize> {
-    match inner {
-        Some(inner) => {
-            let result = std::io::copy(&mut Cursor::new(input), inner).map(|v| v as usize);
+pub(crate) fn stream_compress<W: Write>(encoder: &mut Option<W>, input: &[u8]) -> PyResult<usize> {
+    match encoder {
+        Some(encoder) => {
+            let result = std::io::copy(&mut Cursor::new(input), encoder).map(|v| v as usize);
             crate::to_py_err!(CompressionError -> result)
         }
         None => Err(CompressionError::new_err(
@@ -612,3 +612,24 @@ pub(crate) fn stream_compress<W: Write>(inner: &mut Option<W>, input: &[u8]) ->
         )),
     }
 }
+
+// general stream finish interface. Can't use associated types due to pyo3::pyclass
+// not supporting generic structs.
+#[inline(always)]
+pub(crate) fn stream_finish<W, F, E>(encoder: &mut Option<W>, into_vec: F) -> PyResult<RustyBuffer>
+where
+    W: Write,
+    E: ToString,
+    F: Fn(W) -> Result<Vec<u8>, E>,
+{
+    let mut detached_encoder = None;
+    std::mem::swap(&mut detached_encoder, encoder);
+
+    match detached_encoder {
+        Some(encoder) => {
+            let result = crate::to_py_err!(CompressionError -> into_vec(encoder))?;
+            Ok(RustyBuffer::from(result))
+        }
+        None => Ok(RustyBuffer::from(vec![])),
+    }
+}

src/lz4.rs

@@ -146,17 +146,10 @@ impl Compressor {
     /// 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> {
-        let mut inner = None;
-        std::mem::swap(&mut inner, &mut self.inner);
-
-        match inner {
-            Some(inner) => {
-                let (w, result) = inner.finish();
-                to_py_err!(CompressionError -> result)?;
-                Ok(RustyBuffer::from(w.into_inner()))
-            }
-            None => Ok(RustyBuffer::from(vec![])),
-        }
+        crate::io::stream_finish(&mut self.inner, |inner| {
+            let (cursor, result) = inner.finish();
+            result.map(|_| cursor.into_inner())
+        })
     }
 }
 

src/snappy.rs

@@ -147,16 +147,7 @@ impl Compressor {
     /// 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> {
-        let mut inner = None;
-        std::mem::swap(&mut inner, &mut self.inner);
-
-        match inner {
-            Some(inner) => {
-                let result = to_py_err!(DecompressionError -> inner.into_inner())?;
-                Ok(RustyBuffer::from(result.into_inner()))
-            }
-            None => Ok(RustyBuffer::from(vec![])),
-        }
+        crate::io::stream_finish(&mut self.inner, |inner| inner.into_inner().map(|c| c.into_inner()))
     }
 }
 

src/zstd.rs

@@ -79,16 +79,7 @@ impl Compressor {
     /// 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> {
-        let mut inner = None;
-        std::mem::swap(&mut inner, &mut self.inner);
-
-        match inner {
-            Some(inner) => {
-                let result = to_py_err!(CompressionError -> inner.finish())?;
-                Ok(RustyBuffer::from(result.into_inner()))
-            }
-            None => Ok(RustyBuffer::from(vec![])),
-        }
+        crate::io::stream_finish(&mut self.inner, |inner| inner.finish().map(|v| v.into_inner()))
     }
 }