feat: refactor archive type

crates/rattler_package_streaming/src/archive.rs

@@ -0,0 +1,53 @@
+//! Functions that enable extracting or streaming a Conda package for objects that implement the
+
+use std::path::{Path, PathBuf};
+
+use rattler_conda_types::package::ArchiveType;
+
+use crate::read::{folder_from_conda, folder_from_tar_bz2};
+/// Test
+pub struct Archive {
+    /// Test
+    pub archive_type: ArchiveType,
+    /// Test
+    pub location: PathBuf,
+}
+
+impl Archive {
+    /// Test
+    pub fn new(archive_type: ArchiveType, location: PathBuf) -> Self {
+        Archive {
+            archive_type,
+            location,
+        }
+    }
+
+    /// Test
+    pub fn extract_a_folder(
+        &self,
+        folder_to_extract: &Path,
+        destination: &Path,
+    ) -> Result<(), std::io::Error> {
+        match self.archive_type {
+            ArchiveType::TarBz2 => {
+                folder_from_tar_bz2(&self.location, folder_to_extract, destination)
+            }
+            ArchiveType::Conda => folder_from_conda(&self.location, folder_to_extract, destination),
+        }
+    }
+}
+
+impl TryFrom<PathBuf> for Archive {
+    type Error = std::io::Error;
+
+    fn try_from(path: PathBuf) -> Result<Self, Self::Error> {
+        let archive_type = ArchiveType::try_from(path.as_path()).ok_or(std::io::Error::new(
+            std::io::ErrorKind::NotFound,
+            "package does not point to valid archive",
+        ))?;
+        Ok(Archive {
+            archive_type,
+            location: path,
+        })
+    }
+}

crates/rattler_package_streaming/src/lib.rs

@@ -10,6 +10,7 @@ use rattler_digest::{Md5Hash, Sha256Hash};
 #[cfg(feature = "reqwest")]
 use rattler_networking::Redact;
 
+pub mod archive;
 pub mod read;
 pub mod seek;
 

crates/rattler_package_streaming/src/read.rs

@@ -1,6 +1,8 @@
 //! Functions that enable extracting or streaming a Conda package for objects that implement the
 //! [`std::io::Read`] trait.
 
+use crate::{read, seek};
+
 use super::{ExtractError, ExtractResult};
 use std::mem::ManuallyDrop;
 use std::{ffi::OsStr, io::Read, path::Path};
@@ -86,3 +88,60 @@ pub fn extract_conda(reader: impl Read, destination: &Path) -> Result<ExtractRes
 
     Ok(ExtractResult { sha256, md5 })
 }
+
+/// Extracts a folder from a tar.bz2 archive.
+pub fn folder_from_tar_bz2(
+    archive_path: &Path,
+    find_path: &Path,
+    dest_folder: &Path,
+) -> Result<(), std::io::Error> {
+    let reader = std::fs::File::open(archive_path)?;
+    let mut archive = read::stream_tar_bz2(reader);
+
+    for entry in archive.entries()? {
+        let mut entry = entry?;
+        let path = entry.path()?;
+        if let Ok(stripped_path) = path.strip_prefix(find_path) {
+            let dest_file = dest_folder.join(stripped_path);
+            if let Some(parent_folder) = dest_file.parent() {
+                if !parent_folder.exists() {
+                    std::fs::create_dir_all(parent_folder)?;
+                }
+            }
+            let mut dest_file = std::fs::File::create(dest_file)?;
+            std::io::copy(&mut entry, &mut dest_file)?;
+        }
+    }
+    Ok(())
+}
+
+/// Extracts a folder from a conda archive.
+pub fn folder_from_conda(
+    archive_path: &Path,
+    find_path: &Path,
+    dest_folder: &Path,
+) -> Result<(), std::io::Error> {
+    let reader = std::fs::File::open(archive_path)?;
+
+    let mut archive = if find_path.starts_with("info") {
+        seek::stream_conda_info(reader).expect("Could not open conda file")
+    } else {
+        todo!("Not implemented yet");
+    };
+
+    for entry in archive.entries()? {
+        let mut entry = entry?;
+        let path = entry.path()?;
+        if let Ok(stripped_path) = path.strip_prefix(find_path) {
+            let dest_file = dest_folder.join(stripped_path);
+            if let Some(parent_folder) = dest_file.parent() {
+                if !parent_folder.exists() {
+                    std::fs::create_dir_all(parent_folder)?;
+                }
+            }
+            let mut dest_file = std::fs::File::create(dest_file)?;
+            std::io::copy(&mut entry, &mut dest_file)?;
+        }
+    }
+    Ok(())
+}