Add support for reading archives

Author: DoctorJohn

pyproject.toml

@@ -15,6 +15,7 @@ build-backend = "maturin"
 [dependency-groups]
 dev = [
     "maturin>=1.9.6",
+    "psutil>=7.1.2",
     "pytest>=8.3.5",
     "ruff>=0.14.2",
 ]

src/lib.rs

@@ -5,19 +5,10 @@ use pyo3::exceptions::PyRuntimeError;
 use pyo3::prelude::*;
 use pyo3::types::{PyAny, PyType};
 use std::fs::File;
-use std::io::Write;
+use std::io::{Read, Write};
 use std::path::{Path, PathBuf};
 use tar::Archive;
 
-#[pyfunction]
-fn untar_gz(tar_gz_path: String, destination_path: String) -> PyResult<()> {
-    let tar_gz = File::open(tar_gz_path)?;
-    let tar = GzDecoder::new(tar_gz);
-    let mut archive = Archive::new(tar);
-    archive.unpack(destination_path)?;
-    Ok(())
-}
-
 #[pyclass]
 struct ArchiveWriter {
     builder: Option<tar::Builder<Box<dyn Write + Send + Sync>>>,
@@ -118,6 +109,81 @@ impl ArchiveWriter {
     }
 }
 
+#[pyclass(unsendable)]
+struct ArchiveReader {
+    archive: Option<Archive<Box<dyn Read>>>,
+}
+
+#[pymethods]
+impl ArchiveReader {
+    #[classmethod]
+    #[pyo3(signature = (path, mode="r:gz"))]
+    fn open(
+        _cls: &Bound<'_, PyType>,
+        py: Python<'_>,
+        path: PathBuf,
+        mode: &str,
+    ) -> PyResult<Py<ArchiveReader>> {
+        match mode {
+            "r:gz" => {
+                let file = File::open(path)?;
+                let decoder = GzDecoder::new(file);
+                let reader: Box<dyn Read> = Box::new(decoder);
+                let archive = Archive::new(reader);
+                Py::new(
+                    py,
+                    ArchiveReader {
+                        archive: Some(archive),
+                    },
+                )
+            }
+            "r" => {
+                let file = File::open(path)?;
+                let reader: Box<dyn Read> = Box::new(file);
+                let archive = Archive::new(reader);
+                Py::new(
+                    py,
+                    ArchiveReader {
+                        archive: Some(archive),
+                    },
+                )
+            }
+            _ => Err(PyRuntimeError::new_err(
+                "unsupported mode; only 'r' and 'r:gz' are supported",
+            )),
+        }
+    }
+
+    fn extract(&mut self, to: PathBuf) -> PyResult<()> {
+        let archive = self
+            .archive
+            .as_mut()
+            .ok_or_else(|| PyRuntimeError::new_err("archive is already closed"))?;
+
+        archive.unpack(to)?;
+        Ok(())
+    }
+
+    fn close(&mut self) -> PyResult<()> {
+        self.archive.take();
+        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> {
@@ -126,16 +192,20 @@ fn open(py: Python<'_>, path: PathBuf, mode: &str) -> PyResult<PyObject> {
             let writer = ArchiveWriter::open(&py.get_type::<ArchiveWriter>(), py, path, mode)?;
             Ok(writer.into())
         }
+        "r" | "r:gz" => {
+            let reader = ArchiveReader::open(&py.get_type::<ArchiveReader>(), py, path, mode)?;
+            Ok(reader.into())
+        }
         _ => Err(PyRuntimeError::new_err(
-            "unsupported mode; supported modes are 'w', 'w:gz'",
+            "unsupported mode; supported modes are 'w', 'w:gz', 'r', 'r:gz'",
         )),
     }
 }
 
 #[pymodule]
 fn fastar(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<ArchiveWriter>()?;
+    m.add_class::<ArchiveReader>()?;
     m.add_function(wrap_pyfunction!(open, m)?)?;
-    m.add_function(wrap_pyfunction!(untar_gz, m)?)?;
     Ok(())
 }

tests/conftest.py

@@ -12,6 +12,20 @@ def archive_path(tmp_path) -> Path:
     return tmp_path / "archive.tar.gz"
 
 
+@pytest.fixture
+def source_path(tmp_path) -> Path:
+    path = tmp_path / "source"
+    path.mkdir()
+    return path
+
+
+@pytest.fixture
+def target_path(tmp_path) -> Path:
+    path = tmp_path / "target"
+    path.mkdir()
+    return path
+
+
 @pytest.fixture(params=[("w", "r"), ("w:gz", "r:gz")])
 def modes(request) -> tuple[WriteMode, ReadMode]:
     return request.param