diff --git a/src/frame/compress.rs b/src/frame/compress.rs index bc2c9d4..8b4f23f 100644 --- a/src/frame/compress.rs +++ b/src/frame/compress.rs @@ -111,6 +111,17 @@ impl FrameEncoder { ); } + /// 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 { + 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 io::Write for FrameEncoder { } } +/// A wrapper around an `FrameEncoder` 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 { + // We wrap this in an option to take it during drop. + encoder: Option>, +} + +impl Drop for AutoFinishEncoder { + 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 Write for AutoFinishEncoder { + fn write(&mut self, buf: &[u8]) -> io::Result { + self.encoder.as_mut().unwrap().write(buf) + } + + fn flush(&mut self) -> io::Result<()> { + self.encoder.as_mut().unwrap().flush() + } +} + impl fmt::Debug for FrameEncoder { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { f.debug_struct("FrameEncoder")