use anyhow::{Context, Result};
use indicatif::{ProgressBar, ProgressStyle};
use std::collections::HashMap;
use std::fs::File;
use std::path::{Path, PathBuf};
use walkdir::WalkDir;
pub fn decompress(archive_path: impl AsRef<Path>, output_dir: impl AsRef<Path>) -> Result<()> {
let archive_path = archive_path.as_ref();
let output_dir = output_dir.as_ref();
std::fs::create_dir_all(output_dir)
.with_context(|| format!("创建输出目录失败: {}", output_dir.display()))?;
let archive_size = std::fs::metadata(archive_path)
.with_context(|| format!("获取压缩文件元数据失败: {}", archive_path.display()))?
.len();
let pb = ProgressBar::new(archive_size);
pb.set_style(
ProgressStyle::default_bar()
.template(
"{spinner:.green} [{elapsed_precise}] [{bar:40.cyan/blue}] {bytes}/{total_bytes} ({bytes_per_sec}) {eta} {msg}",
)
.unwrap()
.progress_chars("#>-"),
);
pb.set_message("准备解压...");
let file = File::open(archive_path)
.with_context(|| format!("无法打开压缩文件: {}", archive_path.display()))?;
let decoder = flate2::read::GzDecoder::new(file);
let mut archive = tar::Archive::new(decoder);
decompress_with_progress(&mut archive, output_dir, &pb)
.with_context(|| format!("解压文件失败: {}", archive_path.display()))?;
pb.finish_with_message("解压完成");
Ok(())
}
pub fn compress(source_dir: impl AsRef<Path>, output_path: impl AsRef<Path>) -> Result<()> {
let source_dir = source_dir.as_ref();
let output_path = output_path.as_ref();
if !source_dir.exists() || !source_dir.is_dir() {
anyhow::bail!("源目录不存在: {}", source_dir.display());
}
let total_bytes = count_total_bytes_in_directory(source_dir)?;
let pb = ProgressBar::new(total_bytes);
pb.set_style(
ProgressStyle::default_bar()
.template(
"{spinner:.green} [{elapsed_precise}] [{bar:40.cyan/blue}] {bytes}/{total_bytes} ({bytes_per_sec}) {eta} {msg}",
)
.unwrap()
.progress_chars("#>-"),
);
pb.set_message("准备压缩...");
let file = File::create(output_path)
.with_context(|| format!("创建输出文件失败: {}", output_path.display()))?;
let encoder = flate2::write::GzEncoder::new(file, flate2::Compression::default());
let mut tar = tar::Builder::new(encoder);
compress_with_progress(&mut tar, source_dir, &pb, &mut HashMap::new())
.with_context(|| format!("添加目录到压缩文件失败: {}", source_dir.display()))?;
pb.finish_with_message("压缩完成");
tar.finish().with_context(|| "完成压缩文件写入失败")?;
Ok(())
}
fn count_total_bytes_in_directory(dir: &Path) -> Result<u64> {
let mut total_bytes = 0;
for entry in WalkDir::new(dir)
.into_iter()
.filter_entry(|e| !is_hidden(e))
{
let entry = entry.with_context(|| "遍历目录失败")?;
let path = entry.path();
if path == dir {
continue;
}
if entry.file_type().is_file() {
let metadata = entry
.metadata()
.with_context(|| format!("获取文件元数据失败: {}", path.display()))?;
total_bytes += metadata.len();
} else if entry.file_type().is_dir() {
total_bytes += 1024;
}
}
Ok(total_bytes)
}
fn is_hidden(entry: &walkdir::DirEntry) -> bool {
entry
.file_name()
.to_str()
.map(|s| s.starts_with('.'))
.unwrap_or(false)
}
fn compress_with_progress(
tar: &mut tar::Builder<flate2::write::GzEncoder<File>>,
source_dir: &Path,
pb: &ProgressBar,
visited: &mut HashMap<PathBuf, bool>,
) -> Result<()> {
for entry in WalkDir::new(source_dir)
.into_iter()
.filter_entry(|e| !is_hidden(e))
{
let entry = entry.with_context(|| "遍历目录失败")?;
let path = entry.path();
if path == source_dir {
continue;
}
if visited.contains_key(path) {
continue;
}
visited.insert(path.to_path_buf(), true);
let relative_path = path
.strip_prefix(source_dir)
.with_context(|| format!("计算相对路径失败: {}", path.display()))?;
if entry.file_type().is_file() {
let mut file =
File::open(path).with_context(|| format!("无法打开文件: {}", path.display()))?;
let metadata = file
.metadata()
.with_context(|| format!("获取文件元数据失败: {}", path.display()))?;
let file_size = metadata.len();
pb.set_message(format!(
"正在压缩: {} ({})",
relative_path.display(),
format_bytes(file_size)
));
tar.append_file(relative_path, &mut file)
.with_context(|| format!("添加文件失败: {}", path.display()))?;
pb.inc(file_size);
} else if entry.file_type().is_dir() {
tar.append_dir(relative_path, path)
.with_context(|| format!("添加目录失败: {}", path.display()))?;
pb.inc(1024);
}
}
Ok(())
}
fn format_bytes(bytes: u64) -> String {
const UNITS: &[&str] = &["B", "KB", "MB", "GB", "TB"];
let mut bytes = bytes as f64;
let mut unit_index = 0;
while bytes >= 1024.0 && unit_index < UNITS.len() - 1 {
bytes /= 1024.0;
unit_index += 1;
}
if unit_index == 0 {
format!("{} {}", bytes as u64, UNITS[unit_index])
} else {
format!("{:.1} {}", bytes, UNITS[unit_index])
}
}
fn decompress_with_progress<R: std::io::Read>(
archive: &mut tar::Archive<R>,
output_dir: &Path,
pb: &ProgressBar,
) -> Result<()> {
let entries = archive.entries().with_context(|| "无法读取归档条目")?;
for entry in entries {
let mut entry = entry.with_context(|| "无法读取归档条目")?;
let path_display = {
let path = entry.path().with_context(|| "无法获取条目路径")?;
path.display().to_string()
};
let file_size = entry.size();
pb.set_message(format!(
"解压: {} ({})",
path_display,
format_bytes(file_size)
));
entry
.unpack_in(output_dir)
.with_context(|| format!("解压条目失败: {path_display}"))?;
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use std::fs;
use tempfile::tempdir;
#[test]
fn test_compress_decompress_roundtrip() {
let temp_dir = tempdir().unwrap();
let source_dir = temp_dir.path().join("source");
let archive_path = temp_dir.path().join("test.tar.gz");
let extract_dir = temp_dir.path().join("extract");
fs::create_dir_all(&source_dir).unwrap();
fs::write(source_dir.join("test.txt"), "Hello, World!").unwrap();
fs::create_dir_all(source_dir.join("subdir")).unwrap();
fs::write(source_dir.join("subdir").join("nested.txt"), "Nested file").unwrap();
compress(&source_dir, &archive_path).unwrap();
assert!(archive_path.exists());
decompress(&archive_path, &extract_dir).unwrap();
assert!(extract_dir.join("test.txt").exists());
assert!(extract_dir.join("subdir").join("nested.txt").exists());
assert_eq!(
fs::read_to_string(extract_dir.join("test.txt")).unwrap(),
"Hello, World!"
);
assert_eq!(
fs::read_to_string(extract_dir.join("subdir").join("nested.txt")).unwrap(),
"Nested file"
);
}
#[test]
fn test_decompress_nonexistent_file() {
let temp_dir = tempdir().unwrap();
let result = decompress("nonexistent.tar.gz", temp_dir.path());
assert!(result.is_err());
}
#[test]
fn test_compress_nonexistent_directory() {
let temp_dir = tempdir().unwrap();
let archive_path = temp_dir.path().join("test.tar.gz");
let result = compress("nonexistent_directory", &archive_path);
assert!(result.is_err());
}
#[test]
fn test_compress_empty_directory() {
let temp_dir = tempdir().unwrap();
let source_dir = temp_dir.path().join("empty");
let archive_path = temp_dir.path().join("empty.tar.gz");
let extract_dir = temp_dir.path().join("extract_empty");
fs::create_dir_all(&source_dir).unwrap();
compress(&source_dir, &archive_path).unwrap();
assert!(archive_path.exists());
decompress(&archive_path, &extract_dir).unwrap();
assert!(extract_dir.exists());
}
#[test]
fn test_compress_with_special_files() {
let temp_dir = tempdir().unwrap();
let source_dir = temp_dir.path().join("special");
let archive_path = temp_dir.path().join("special.tar.gz");
let extract_dir = temp_dir.path().join("extract_special");
fs::create_dir_all(&source_dir).unwrap();
fs::write(source_dir.join("normal.txt"), "Normal file").unwrap();
fs::create_dir_all(source_dir.join("subdir")).unwrap();
fs::write(source_dir.join("subdir").join("file.txt"), "Subdir file").unwrap();
compress(&source_dir, &archive_path).unwrap();
assert!(archive_path.exists());
decompress(&archive_path, &extract_dir).unwrap();
assert!(extract_dir.join("normal.txt").exists());
assert!(extract_dir.join("subdir").exists());
assert!(extract_dir.join("subdir").join("file.txt").exists());
assert_eq!(
fs::read_to_string(extract_dir.join("normal.txt")).unwrap(),
"Normal file"
);
assert_eq!(
fs::read_to_string(extract_dir.join("subdir").join("file.txt")).unwrap(),
"Subdir file"
);
}
#[test]
fn test_decompress_to_existing_directory() {
let temp_dir = tempdir().unwrap();
let source_dir = temp_dir.path().join("source");
let archive_path = temp_dir.path().join("test.tar.gz");
let extract_dir = temp_dir.path().join("existing");
fs::create_dir_all(&source_dir).unwrap();
fs::write(source_dir.join("test.txt"), "Test content").unwrap();
fs::create_dir_all(&extract_dir).unwrap();
fs::write(extract_dir.join("existing.txt"), "Existing content").unwrap();
compress(&source_dir, &archive_path).unwrap();
decompress(&archive_path, &extract_dir).unwrap();
assert!(extract_dir.join("test.txt").exists());
assert!(extract_dir.join("existing.txt").exists());
assert_eq!(
fs::read_to_string(extract_dir.join("test.txt")).unwrap(),
"Test content"
);
assert_eq!(
fs::read_to_string(extract_dir.join("existing.txt")).unwrap(),
"Existing content"
);
}
}