rapiddoc_core/
progress.rs

1//! Shared tqdm-style presentation for terminals, IDEs and redirected stderr.
2use std::collections::BTreeMap;
3use std::io::{self, IsTerminal, Write};
4use std::time::{Duration, Instant};
5pub fn stage_name(stage: &str) -> &str {
6    match stage {
7        "输入分类" => "classify",
8        "原生全文分析" => "native_analysis",
9        "原生分析失败，回退版面模型" => "layout_fallback",
10        "读取文字与渲染" => "render",
11        "版面模型识别" => "layout",
12        "公式识别" => "formula",
13        "正文 OCR 检测与原生文字规划" => "ocr_detection",
14        "正文 OCR 批量识别" => "ocr_recognition",
15        "表格识别" => "table",
16        "印章识别" => "seal",
17        "页面组装与图片资产" => "assemble",
18        "补识别、段落和跨页表格" => "finalize",
19        "写出结果" => "write",
20        _ => stage,
21    }
22}
23fn label(stage: &str) -> &str {
24    match stage {
25        "layout" => "Layout Predict",
26        "formula" => "Formula Predict",
27        "ocr_detection" => "PDF/OCR-det Predict",
28        "ocr_recognition" => "OCR-rec Predict",
29        "table" => "Table Predict",
30        "assemble" => "Processing pages",
31        "orientation" => "Orientation Predict",
32        "seal" => "Seal Predict",
33        "finalize" => "Post-processing",
34        _ => stage,
35    }
36}
37fn clock(seconds: f64) -> String {
38    let s = seconds.max(0.) as u64;
39    if s >= 3600 {
40        format!("{}:{:02}:{:02}", s / 3600, (s / 60) % 60, s % 60)
41    } else {
42        format!("{:02}:{:02}", s / 60, s % 60)
43    }
44}
45pub fn format_stage(stage: &str, done: usize, total: usize, elapsed: Duration) -> String {
46    let done = done.min(total);
47    let fill = if total == 0 {
48        20
49    } else {
50        done.saturating_mul(20) / total
51    };
52    let percent = if total == 0 {
53        100
54    } else {
55        done.saturating_mul(100) / total
56    };
57    let seconds = elapsed.as_secs_f64();
58    let (eta, rate) = if done > 0 && seconds > 0. {
59        let rate = done as f64 / seconds;
60        (
61            clock((total - done) as f64 / rate),
62            if rate >= 1. {
63                format!("{rate:.2}it/s")
64            } else {
65                format!("{:.2}s/it", 1. / rate)
66            },
67        )
68    } else if total == 0 {
69        ("00:00".into(), "0.00it/s".into())
70    } else {
71        ("?".into(), "?it/s".into())
72    };
73    format!(
74        "{}: {:3}%|{}{}| {}/{} [{}<{}, {}]",
75        label(stage),
76        percent,
77        "█".repeat(fill),
78        " ".repeat(20 - fill),
79        done,
80        total,
81        clock(seconds),
82        eta,
83        rate
84    )
85}
86struct Stage {
87    measured: Option<Duration>,
88    done: usize,
89    total: usize,
90    started: Instant,
91    finished: bool,
92}
93pub struct TerminalProgress {
94    enabled: bool,
95    dynamic: bool,
96    stages: BTreeMap<String, Stage>,
97    last_draw: Option<Instant>,
98    width: usize,
99    active: Option<String>,
100}
101impl TerminalProgress {
102    pub fn new(_name: String, _total: usize) -> Self {
103        let mode = std::env::var("RAPIDDOC_PROGRESS")
104            .unwrap_or_default()
105            .to_ascii_lowercase();
106        let dynamic = match mode.as_str() {
107            "always" | "1" => true,
108            "plain" => false,
109            _ => io::stderr().is_terminal() || std::env::var_os("PYCHARM_HOSTED").is_some(),
110        };
111        Self {
112            enabled: !matches!(mode.as_str(), "off" | "0"),
113            dynamic,
114            stages: BTreeMap::new(),
115            last_draw: None,
116            width: 0,
117            active: None,
118        }
119    }
120    pub fn update(&mut self, stage: &str, page: Option<usize>) {
121        self.update_detail(stage, page, None)
122    }
123    pub fn update_detail(&mut self, stage: &str, _page: Option<usize>, _detail: Option<&str>) {
124        if matches!(stage_name(stage), "finalize" | "write") {
125            self.clear();
126            if self.enabled {
127                let _ = writeln!(io::stderr().lock(), "{}...", label(stage));
128            }
129        }
130    }
131    pub fn reset(&mut self) {
132        self.clear();
133        self.stages.clear();
134        self.active = None;
135    }
136    pub fn advance(&mut self, stage: &str, done: usize, total: usize) {
137        self.advance_with_elapsed(stage, done, total, None)
138    }
139    /// Production events carry measured stage work; gaps between pages are not work.
140    pub fn advance_with_elapsed(
141        &mut self,
142        stage: &str,
143        done: usize,
144        total: usize,
145        elapsed: Option<Duration>,
146    ) {
147        if !self.enabled || total == 0 {
148            return;
149        }
150        let stage = stage_name(stage);
151        let fresh = !self.stages.contains_key(stage);
152        let state = self.stages.entry(stage.into()).or_insert_with(|| Stage {
153            measured: None,
154            done: 0,
155            total,
156            started: Instant::now(),
157            finished: false,
158        });
159        if state.finished {
160            return;
161        }
162        state.done = state.done.max(done.min(total));
163        state.total = total;
164        if let Some(elapsed) = elapsed {
165            state.measured = Some(state.measured.unwrap_or_default().max(elapsed));
166        }
167        let finished = state.done == total;
168        let due = self.last_draw.is_none_or(|t| {
169            t.elapsed()
170                >= if self.dynamic {
171                    Duration::from_millis(100)
172                } else {
173                    Duration::from_secs(5)
174                }
175        });
176        let switched = self.active.as_deref() != Some(stage);
177        if !finished && !fresh && !due && !(self.dynamic && switched) {
178            return;
179        }
180        let text = format_stage(
181            stage,
182            state.done,
183            total,
184            state.measured.unwrap_or_else(|| state.started.elapsed()),
185        );
186        state.finished = finished;
187        let mut stderr = io::stderr().lock();
188        if self.dynamic {
189            let _ = write!(
190                stderr,
191                "\r{}{}",
192                text,
193                " ".repeat(self.width.saturating_sub(text.chars().count()))
194            );
195            if finished {
196                let _ = writeln!(stderr);
197                self.width = 0;
198            } else {
199                self.width = text.chars().count();
200            }
201        } else {
202            let _ = writeln!(stderr, "{text}");
203        }
204        let _ = stderr.flush();
205        self.last_draw = Some(Instant::now());
206        self.active = Some(stage.into());
207    }
208    pub fn clear(&mut self) {
209        if self.dynamic && self.width > 0 {
210            let mut stderr = io::stderr().lock();
211            let _ = write!(stderr, "\r{}\r", " ".repeat(self.width));
212            let _ = stderr.flush();
213            self.width = 0;
214        }
215        self.last_draw = None;
216    }
217}
218impl Drop for TerminalProgress {
219    fn drop(&mut self) {
220        self.clear();
221    }
222}
223#[cfg(test)]
224mod tests {
225    use super::*;
226    #[test]
227    fn stage_counts_eta_rate_and_zero_work() {
228        let line = format_stage("formula", 2, 5, Duration::from_secs(4));
229        assert!(line.starts_with("Formula Predict:  40%|"));
230        assert!(line.ends_with("2/5 [00:04<00:06, 2.00s/it]"));
231        assert!(format_stage("layout", 14, 14, Duration::from_secs(7))
232            .ends_with("14/14 [00:07<00:00, 2.00it/s]"));
233        assert!(format_stage("table", 0, 0, Duration::ZERO).contains("100%"));
234        assert!(format_stage("formula", 0, 5, Duration::ZERO).contains("00:00<?, ?it/s"));
235    }
236}
rapiddoc_core/progress.rs

rapiddoc_core/
progress.rs
