diff --git a/tsar-py/src/lib.rs b/tsar-py/src/lib.rs index e9d1517..2144d2e 100644 --- a/tsar-py/src/lib.rs +++ b/tsar-py/src/lib.rs @@ -11,7 +11,7 @@ use rayon::prelude::*; #[pyclass(module = "tsar.tsar")] struct Writer { - w: tsar::Builder, + w: Option>, } #[pymethods] @@ -19,12 +19,16 @@ impl Writer { #[new] fn new(dst: &str) -> PyResult { Ok(Self { - w: tsar::Builder::new(std::fs::File::create(dst)?), + w: Some(tsar::Builder::new(std::fs::File::create(dst)?)), }) } pub fn write_file(&mut self, name: String, d: &[u8]) -> PyResult<()> { - self.w.add_file(name, std::io::Cursor::new(d)).unwrap(); + self.w + .as_mut() + .expect("The writer has already been closed.") + .add_file(name, std::io::Cursor::new(d)) + .unwrap(); Ok(()) } @@ -57,18 +61,25 @@ impl Writer { _ => None, }; + let w = self + .w + .as_mut() + .expect("The writer has already been closed."); + match ty { - Some(ty) => self.w.add_blob(name, data, ty, &dims, opt), - _ => self - .w - .add_blob(name, data, tsar::DataType::Byte, &[data.len()], opt), + Some(ty) => w.add_blob(name, data, ty, &dims, opt), + _ => w.add_blob(name, data, tsar::DataType::Byte, &[data.len()], opt), } .unwrap(); Ok(()) } pub fn close(&mut self) -> PyResult<()> { - self.w.finish().unwrap(); + self.w + .take() + .expect("The writer has already been closed.") + .finish() + .unwrap(); Ok(()) } } diff --git a/tsar-rs/src/write/mod.rs b/tsar-rs/src/write/mod.rs index 6441ef4..ee87f12 100644 --- a/tsar-rs/src/write/mod.rs +++ b/tsar-rs/src/write/mod.rs @@ -112,15 +112,17 @@ impl Builder { self.write_chunks(b, [data]) } - pub fn finish(&mut self) -> Result<()> { - self.z - .start_file(paths::BUNDLE_META_PATH, SimpleFileOptions::default())?; - self.meta.blobs.sort_by(|a, b| a.name.cmp(&b.name)); + pub fn finish(self) -> Result<()> { + let Self { + mut z, + mut meta, + chunks: _, + } = self; + z.start_file(paths::BUNDLE_META_PATH, SimpleFileOptions::default())?; + meta.blobs.sort_by(|a, b| a.name.cmp(&b.name)); // TODO check target_file contiguous - self.meta - .write_to(&mut CodedOutputStream::new(&mut self.z)) - .unwrap(); - self.z.finish()?; + meta.write_to(&mut CodedOutputStream::new(&mut z)).unwrap(); + z.finish()?; Ok(()) }