use std::io::{IsTerminal, stderr};
use std::time::{Duration, Instant};
const REDRAW_INTERVAL: Duration = Duration::from_millis(100);
const BAR_WIDTH: usize = 40;
const FILE_NAME_WIDTH: usize = 48;
const LINE_WIDTH: usize = 128;
pub struct Progress {
label: &'static str,
file: String,
total: u64,
done: u64,
last_render: Option<Instant>,
tty: bool,
finished: bool,
}
impl Progress {
pub fn start(label: &'static str, total_bytes: u64) -> Self {
Self {
label,
file: String::new(),
total: total_bytes,
done: 0,
last_render: None,
tty: stderr().is_terminal(),
finished: false,
}
}
pub fn set_file(&mut self, name: &str) {
self.file = tail(name, FILE_NAME_WIDTH);
self.draw();
}
pub fn add(&mut self, bytes: u64) {
self.done += bytes;
let throttled = self
.last_render
.is_some_and(|t| t.elapsed() < REDRAW_INTERVAL);
if !throttled {
self.draw();
}
}
pub fn finish(&mut self) {
if !self.finished {
self.finished = true;
self.draw();
if self.tty {
eprintln!();
}
}
}
fn draw(&mut self) {
if !self.tty || self.finished {
return;
}
self.last_render = Some(Instant::now());
let pct = if self.total > 0 {
(self.done.min(self.total)) as f64 / self.total as f64 * 100.0
} else {
100.0
};
let filled = ((pct / 100.0) * BAR_WIDTH as f64).round() as usize;
let filled = filled.min(BAR_WIDTH);
let (head, rest) = match filled {
0 => (String::new(), ".".repeat(BAR_WIDTH)),
BAR_WIDTH => ("=".repeat(BAR_WIDTH), String::new()),
_ => (
format!("{}>", "=".repeat(filled - 1)),
".".repeat(BAR_WIDTH - filled),
),
};
let msg = format!(
"{} [{head}{rest}] {:>5.1}% {:>9}/{:<9} {}",
self.label,
pct,
fmt_bytes(self.done),
fmt_bytes(self.total),
self.file,
);
let msg = truncate(&msg, LINE_WIDTH);
eprint!("\r{msg:<LINE_WIDTH$}");
}
}
pub fn stderr_note(msg: &str) {
if stderr().is_terminal() {
eprintln!("{msg}");
}
}
pub fn fmt_bytes(n: u64) -> String {
const UNITS: [&str; 5] = ["B", "KB", "MB", "GB", "TB"];
let mut v = n as f64;
for (i, unit) in UNITS.iter().enumerate() {
if v < 1024.0 || i == UNITS.len() - 1 {
return if i == 0 {
format!("{n}B")
} else {
format!("{v:.1}{unit}")
};
}
v /= 1024.0;
}
unreachable!("units 末项已兜底返回")
}
fn tail(s: &str, max_chars: usize) -> String {
if s.chars().count() <= max_chars {
s.to_string()
} else {
let skip = s.chars().count() - max_chars;
s.chars().skip(skip).collect()
}
}
fn truncate(s: &str, max_chars: usize) -> String {
s.chars().take(max_chars).collect()
}
#[cfg(test)]
mod tests {
use super::{Progress, fmt_bytes};
#[test]
fn test_fmt_bytes() {
assert_eq!(fmt_bytes(0), "0B");
assert_eq!(fmt_bytes(512), "512B");
assert_eq!(fmt_bytes(1536), "1.5KB");
assert_eq!(fmt_bytes(18 * 1024 * 1024 * 1024), "18.0GB");
}
#[test]
fn test_progress_non_tty_silent() {
let mut p = Progress::start("test", 1000);
p.set_file("a.jsonl");
p.add(400);
p.add(600);
p.finish();
p.finish();
}
}