Add support for writing archives (#1)

Author: DoctorJohn

pyproject.toml

@@ -13,7 +13,11 @@ requires = ["maturin>=1.9,<2.0"]
 build-backend = "maturin"
 
 [dependency-groups]
-dev = ["maturin>=1.9.6"]
+dev = [
+    "maturin>=1.9.6",
+    "pytest>=8.3.5",
+    "ruff>=0.14.2",
+]
 
 [tool.maturin]
 features = ["pyo3/extension-module"]

src/lib.rs

@@ -1,6 +1,12 @@
 use flate2::read::GzDecoder;
+use flate2::write::GzEncoder;
+use flate2::Compression;
+use pyo3::exceptions::PyRuntimeError;
 use pyo3::prelude::*;
+use pyo3::types::{PyAny, PyType};
 use std::fs::File;
+use std::io::Write;
+use std::path::{Path, PathBuf};
 use tar::Archive;
 
 #[pyfunction]
@@ -12,8 +18,124 @@ fn untar_gz(tar_gz_path: String, destination_path: String) -> PyResult<()> {
     Ok(())
 }
 
+#[pyclass]
+struct ArchiveWriter {
+    builder: Option<tar::Builder<Box<dyn Write + Send + Sync>>>,
+}
+
+#[pymethods]
+impl ArchiveWriter {
+    #[classmethod]
+    #[pyo3(signature = (path, mode="w:gz"))]
+    fn open(
+        _cls: &Bound<'_, PyType>,
+        py: Python<'_>,
+        path: PathBuf,
+        mode: &str,
+    ) -> PyResult<Py<ArchiveWriter>> {
+        match mode {
+            "w:gz" => {
+                let file = File::create(path)?;
+                let enc = GzEncoder::new(file, Compression::default());
+                let writer: Box<dyn Write + Send + Sync> = Box::new(enc);
+                let builder = tar::Builder::new(writer);
+                Py::new(
+                    py,
+                    ArchiveWriter {
+                        builder: Some(builder),
+                    },
+                )
+            }
+            "w" => {
+                let file = File::create(path)?;
+                let writer: Box<dyn Write + Send + Sync> = Box::new(file);
+                let builder = tar::Builder::new(writer);
+                Py::new(
+                    py,
+                    ArchiveWriter {
+                        builder: Some(builder),
+                    },
+                )
+            }
+            _ => Err(PyRuntimeError::new_err(
+                "unsupported mode; only 'w' and 'w:gz' are supported",
+            )),
+        }
+    }
+
+    #[pyo3(signature = (path, arcname=None, recursive=true))]
+    fn add(&mut self, path: PathBuf, arcname: Option<String>, recursive: bool) -> PyResult<()> {
+        let builder = self
+            .builder
+            .as_mut()
+            .ok_or_else(|| PyRuntimeError::new_err("archive is already closed"))?;
+
+        let default_name = || -> PyResult<String> {
+            let name = Path::new(&path)
+                .file_name()
+                .ok_or_else(|| PyRuntimeError::new_err("cannot derive name from path"))?
+                .to_string_lossy()
+                .into_owned();
+            Ok(name)
+        }()?;
+
+        let name = arcname.unwrap_or(default_name);
+
+        if path.is_dir() {
+            if recursive {
+                builder.append_dir_all(&name, &path)?;
+            } else {
+                builder.append_dir(&name, &path)?;
+            }
+        } else if path.is_file() {
+            builder.append_path_with_name(&path, &name)?;
+        } else {
+            return Err(PyRuntimeError::new_err("path does not exist"));
+        }
+        Ok(())
+    }
+
+    fn close(&mut self) -> PyResult<()> {
+        if let Some(builder) = self.builder.take() {
+            let mut writer = builder.into_inner()?;
+            writer.flush()?;
+        }
+        Ok(())
+    }
+
+    fn __enter__(py_self: PyRef<'_, Self>) -> PyRef<'_, Self> {
+        py_self
+    }
+
+    fn __exit__(
+        &mut self,
+        _exc_type: Option<Bound<'_, PyAny>>,
+        _exc: Option<Bound<'_, PyAny>>,
+        _tb: Option<Bound<'_, PyAny>>,
+    ) -> PyResult<bool> {
+        self.close()?;
+        Ok(false) // Propagate exceptions if any
+    }
+}
+
+#[pyfunction]
+#[pyo3(signature = (path, mode))]
+fn open(py: Python<'_>, path: PathBuf, mode: &str) -> PyResult<PyObject> {
+    match mode {
+        "w" | "w:gz" => {
+            let writer = ArchiveWriter::open(&py.get_type::<ArchiveWriter>(), py, path, mode)?;
+            Ok(writer.into())
+        }
+        _ => Err(PyRuntimeError::new_err(
+            "unsupported mode; supported modes are 'w', 'w:gz'",
+        )),
+    }
+}
+
 #[pymodule]
 fn fastar(m: &Bound<'_, PyModule>) -> PyResult<()> {
+    m.add_class::<ArchiveWriter>()?;
+    m.add_function(wrap_pyfunction!(open, m)?)?;
     m.add_function(wrap_pyfunction!(untar_gz, m)?)?;
     Ok(())
 }

tests/conftest.py

@@ -0,0 +1,27 @@
+from typing import Literal, TypeAlias
+from pathlib import Path
+import pytest
+
+
+WriteMode: TypeAlias = Literal["w", "w:gz"]
+ReadMode: TypeAlias = Literal["r", "r:gz"]
+
+
+@pytest.fixture
+def archive_path(tmp_path) -> Path:
+    return tmp_path / "archive.tar.gz"
+
+
+@pytest.fixture(params=[("w", "r"), ("w:gz", "r:gz")])
+def modes(request) -> tuple[WriteMode, ReadMode]:
+    return request.param
+
+
+@pytest.fixture
+def write_mode(modes) -> WriteMode:
+    return modes[0]
+
+
+@pytest.fixture
+def read_mode(modes) -> ReadMode:
+    return modes[1]