feat: add auto_finish to FrameEncoder

PSeitz · PSeitz · commit ce28ee659a62 · 2023-03-04T22:44:45.000+08:00
Fixes #94
diff --git a/src/frame/compress.rs b/src/frame/compress.rs
@@ -111,6 +111,17 @@ impl<W: io::Write> FrameEncoder<W> {
         );
     }
 
+    /// Returns a wrapper around `self` that will finish the stream on drop.
+    ///
+    /// # Panic
+    ///
+    /// Panics on drop if an error happens when finishing the stream.
+    pub fn auto_finish(self) -> AutoFinishEncoder<W> {
+        AutoFinishEncoder {
+            encoder: Some(self),
+        }
+    }
+
     /// Creates a new Encoder with the specified FrameInfo.
     pub fn with_frame_info(frame_info: FrameInfo, wtr: W) -> Self {
         FrameEncoder {
@@ -376,6 +387,37 @@ impl<W: io::Write> io::Write for FrameEncoder<W> {
     }
 }
 
+/// A wrapper around an `FrameEncoder<W>` that finishes the stream on drop.
+///
+/// This can be created by the [`auto_finish()`] method on the [`FrameEncoder`].
+///
+/// [`auto_finish()`]: Encoder::auto_finish
+/// [`Encoder`]: Encoder
+pub struct AutoFinishEncoder<W: Write> {
+    // We wrap this in an option to take it during drop.
+    encoder: Option<FrameEncoder<W>>,
+}
+
+impl<W: io::Write> Drop for AutoFinishEncoder<W> {
+    fn drop(&mut self) {
+        if let Some(mut encoder) = self.encoder.take() {
+            if let Err(err) = encoder.try_finish() {
+                panic!("Error when flushing frame on drop {:?} ", err);
+            }
+        }
+    }
+}
+
+impl<W: Write> Write for AutoFinishEncoder<W> {
+    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
+        self.encoder.as_mut().unwrap().write(buf)
+    }
+
+    fn flush(&mut self) -> io::Result<()> {
+        self.encoder.as_mut().unwrap().flush()
+    }
+}
+
 impl<W: fmt::Debug + io::Write> fmt::Debug for FrameEncoder<W> {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         f.debug_struct("FrameEncoder")
