Skip to main content

ztr_lib/
compressor.rs

1use crate::config::Config;
2use anyhow::{Context, Result};
3use indicatif::{ProgressBar, ProgressStyle};
4use std::fs::File;
5use std::path::{Path, PathBuf};
6
7/// 根据配置压缩指定文件列表。
8
9/// # 参数
10/// - `config`: 压缩配置。
11/// - `base_dir`: 基础目录,所有文件路径都将相对于此目录进行计算。
12/// - `files_to_compress`: 要压缩的文件路径列表。
13
14/// # 返回
15/// `Result<PathBuf>`: 成功时返回输出文件的路径,失败时返回错误信息。
16pub fn compress_directory(
17    config: &Config,
18    base_dir: &Path,
19    files_to_compress: Vec<PathBuf>,
20) -> Result<PathBuf> {
21    let output_name = config.get_output_name();
22    let output_path = match config.format.as_str() {
23        "zip" => base_dir.join(format!("{}.zip", output_name)),
24        "tar.gz" => base_dir.join(format!("{}.tar.gz", output_name)),
25        "7z" => base_dir.join(format!("{}.7z", output_name)),
26        _ => anyhow::bail!("不支持的压缩格式: {}", config.format),
27    };
28
29    println!("正在压缩目录: {}", base_dir.display());
30    println!("输出文件: {}", output_path.display());
31    println!("压缩格式: {}", config.format);
32
33    let files = files_to_compress;
34
35    if files.is_empty() {
36        println!("没有需要压缩的文件。");
37        return Ok(output_path);
38    }
39
40    println!("找到 {} 个文件要压缩", files.len());
41
42    // 创建进度条
43    let pb = ProgressBar::new(files.len() as u64);
44    pb.set_style(
45        ProgressStyle::default_bar()
46            .template("{spinner:.green} [{elapsed_precise}] [{bar:40.cyan/blue}] {pos}/{len} {msg}")
47            .unwrap()
48            .progress_chars("#>-"),
49    );
50    pb.set_message("正在压缩...");
51
52    // 根据格式选择压缩方法
53    let result = match config.format.as_str() {
54        "zip" => compress_zip(&files, base_dir, &output_path, &pb),
55        "tar.gz" => compress_tar_gz(&files, base_dir, &output_path, &pb),
56        "7z" => compress_7z(&files, base_dir, &output_path, &pb),
57        _ => anyhow::bail!("不支持的压缩格式: {}", config.format),
58    };
59
60    pb.finish_with_message("压缩完成");
61
62    match result {
63        Ok(_) => {
64            println!("✓ 压缩完成: {}", output_path.display());
65
66            // 显示文件大小
67            if let Ok(metadata) = std::fs::metadata(&output_path) {
68                let size = metadata.len();
69                if size > 1024 * 1024 {
70                    println!("文件大小: {:.2} MB", size as f64 / (1024.0 * 1024.0));
71                } else if size > 1024 {
72                    println!("文件大小: {:.2} KB", size as f64 / 1024.0);
73                } else {
74                    println!("文件大小: {} bytes", size);
75                }
76            }
77        }
78        Err(e) => {
79            println!("✗ 压缩失败: {}", e);
80            return Err(e);
81        }
82    }
83
84    Ok(output_path)
85}
86
87/// 将文件压缩为 ZIP 格式。
88
89/// # 参数
90/// - `files`: 要压缩的文件路径列表。
91/// - `base_dir`: 基础目录,用于计算文件中相对路径。
92/// - `output_path`: 输出 ZIP 文件的路径。
93/// - `pb`: 进度条。
94
95/// # 返回
96/// `Result<()>`: 成功时返回 `Ok(())`,失败时返回错误信息。
97pub fn compress_zip(
98    files: &[PathBuf],
99    base_dir: &Path,
100    output_path: &Path,
101    pb: &ProgressBar,
102) -> Result<()> {
103    use std::io::Write;
104    use zip::{ZipWriter, write::FileOptions};
105
106    let file = File::create(output_path).context("创建ZIP文件失败")?;
107    let mut zip = ZipWriter::new(file);
108
109    for file_path in files {
110        pb.inc(1);
111
112        let relative_path = file_path
113            .strip_prefix(base_dir)
114            .with_context(|| format!("计算相对路径失败: {}", file_path.display()))?;
115
116        let mut file = File::open(file_path)
117            .with_context(|| format!("打开文件失败: {}", file_path.display()))?;
118
119        let relative_path_str = relative_path.to_string_lossy().replace("\\", "/");
120        zip.start_file(&relative_path_str, FileOptions::default())
121            .with_context(|| format!("添加文件到ZIP失败: {}", file_path.display()))?;
122
123        let mut buffer = Vec::new();
124        std::io::copy(&mut file, &mut buffer)
125            .with_context(|| format!("读取文件失败: {}", file_path.display()))?;
126
127        zip.write_all(&buffer)
128            .with_context(|| format!("写入ZIP失败: {}", file_path.display()))?;
129    }
130
131    zip.finish().context("完成ZIP写入失败")?;
132
133    Ok(())
134}
135
136/// 将文件压缩为 TAR.GZ 格式。
137
138/// # 参数
139/// - `files`: 要压缩的文件路径列表。
140/// - `base_dir`: 基础目录,用于计算文件中相对路径。
141/// - `output_path`: 输出 TAR.GZ 文件的路径。
142/// - `pb`: 进度条。
143
144/// # 返回
145/// `Result<()>`: 成功时返回 `Ok(())`,失败时返回错误信息。
146pub fn compress_tar_gz(
147    files: &[PathBuf],
148    base_dir: &Path,
149    output_path: &Path,
150    pb: &ProgressBar,
151) -> Result<()> {
152    use flate2::Compression;
153    use flate2::write::GzEncoder;
154    use tar::Builder;
155
156    let file = File::create(output_path).context("创建TAR.GZ文件失败")?;
157    let gz_encoder = GzEncoder::new(file, Compression::default());
158    let mut tar = Builder::new(gz_encoder);
159
160    for file_path in files {
161        pb.inc(1);
162
163        let relative_path = file_path
164            .strip_prefix(base_dir)
165            .with_context(|| format!("计算相对路径失败: {}", file_path.display()))?;
166
167        tar.append_path_with_name(file_path, relative_path)
168            .with_context(|| format!("添加文件到TAR失败: {}", file_path.display()))?;
169    }
170
171    tar.finish().context("完成TAR.GZ写入失败")?;
172
173    Ok(())
174}
175
176/// 将文件压缩为 7Z 格式。
177
178/// # 参数
179/// - `files`: 要压缩的文件路径列表。
180/// - `base_dir`: 基础目录,用于计算文件中相对路径。
181/// - `output_path`: 输出 7Z 文件的路径。
182/// - `pb`: 进度条。
183
184/// # 返回
185/// `Result<()>`: 成功时返回 `Ok(())`,失败时返回错误信息。
186pub fn compress_7z(
187    files: &[PathBuf],
188    base_dir: &Path,
189    output_path: &Path,
190    pb: &ProgressBar,
191) -> Result<()> {
192    use sevenz_rust::{SevenZArchiveEntry, SevenZWriter};
193    use std::io::Read;
194
195    let mut sz_writer = SevenZWriter::create(output_path).context("创建7Z文件失败")?;
196
197    for file_path in files {
198        pb.inc(1);
199
200        let relative_path = file_path
201            .strip_prefix(base_dir)
202            .with_context(|| format!("计算相对路径失败: {}", file_path.display()))?;
203
204        if file_path.is_file() {
205            let mut file = File::open(file_path)
206                .with_context(|| format!("打开文件失败: {}", file_path.display()))?;
207            let mut content = Vec::new();
208            file.read_to_end(&mut content)
209                .with_context(|| format!("读取文件内容失败: {}", file_path.display()))?;
210
211            let mut entry = SevenZArchiveEntry::default();
212            entry.name = relative_path.to_string_lossy().replace("\\", "/");
213            entry.size = content.len() as u64;
214
215            sz_writer
216                .push_archive_entry(entry, Some(content.as_slice()))
217                .with_context(|| format!("添加文件到7Z失败: {}", file_path.display()))?;
218        }
219    }
220
221    sz_writer.finish().context("完成7Z写入失败")?;
222
223    Ok(())
224}