file: 增加 set_download_process 函数,可设置是否显示下载进度

Signed-off-by: Jia Chao <jiac13@chinaunicom.cn>
This commit is contained in:
Jia Chao 2024-07-10 11:59:55 +08:00
parent f5b40e8cfe
commit ee16adb37f

View File

@ -2,12 +2,12 @@ use std::env;
use std::fs;
use std::path::{Path, PathBuf};
use tokio::fs::{File, OpenOptions};
use tokio::io::{AsyncSeekExt, AsyncWriteExt, SeekFrom};
use futures_util::StreamExt;
use indicatif::{ProgressBar, ProgressStyle};
use path_absolutize::Absolutize;
use reqwest::Client;
use tokio::fs::{File, OpenOptions};
use tokio::io::{AsyncSeekExt, AsyncWriteExt, SeekFrom};
use tracing::{debug, info};
use crate::util::uuid;
@ -84,6 +84,32 @@ pub fn walk_dir<P: AsRef<Path>>(path: P, nodir: bool) -> Vec<PathBuf> {
res
}
/// 设置或清除下载进程的环境变量 `CCUTIL_DOWNLOAD_PROCESS`。在使用 `download` 和 `async_download`
/// 函数时,会根据是否设置此变量来显示或隐藏下载进度。
///
/// # 参数
///
/// - `with_process`: 一个布尔值,指示是否要设置下载进程环境变量。
///
/// # 示例
///
/// ```rust
/// // 设置下载进程环境变量
/// set_download_process(true);
///
/// // 清除下载进程环境变量
/// set_download_process(false);
/// ```
pub fn set_download_process(with_process: bool) {
if with_process {
// 设置环境变量为字符串 "1"
env::set_var("CCUTIL_DOWNLOAD_PROCESS", "1");
} else {
// 清除环境变量
env::remove_var("CCUTIL_DOWNLOAD_PROCESS");
}
}
/// 从给定的 URL 下载文件并将其保存到指定的目标路径。
///
/// 该函数初始化一个异步运行时,并阻塞在异步下载函数上,直到下载完成。
@ -208,11 +234,17 @@ pub async fn async_download<P: AsRef<Path>>(url: String, target: Option<P>) -> c
// 设置下载进度条
let pb = ProgressBar::new(total_size);
pb.set_style(
ProgressStyle::default_bar()
.template("{spinner:.green} [{elapsed_precise}] [{wide_bar:.cyan/blue}] {bytes}/{total_bytes} ({bytes_per_sec}, {eta_precise})").unwrap()
.progress_chars("#>-"),
);
let with_process = match env::var("CCUTIL_DOWNLOAD_PROCESS") {
Ok(_) => {
pb.set_style(
ProgressStyle::default_bar()
.template("{spinner:.green} [{elapsed_precise}] [{wide_bar:.cyan/blue}] {bytes}/{total_bytes} ({bytes_per_sec}, {eta_precise})").unwrap()
.progress_chars("#>-"),
);
true
}
Err(_) => false,
};
// 读取响应数据并写入文件
while let Some(chunk) = stream.next().await {
@ -221,7 +253,9 @@ pub async fn async_download<P: AsRef<Path>>(url: String, target: Option<P>) -> c
current_size += chunk.len() as u64;
// 更新下载进度
pb.set_position(current_size);
if with_process {
pb.set_position(current_size);
}
}
Ok(())