edit prediction cli: Cargo-style progress output (#44675)

Release Notes:

- N/A

crates/edit_prediction_cli/src/progress.rs

@@ -0,0 +1,372 @@
+use std::{
+    borrow::Cow,
+    collections::HashMap,
+    io::{IsTerminal, Write},
+    sync::{Arc, Mutex},
+    time::{Duration, Instant},
+};
+
+pub struct Progress {
+    inner: Mutex<ProgressInner>,
+}
+
+struct ProgressInner {
+    completed: Vec<CompletedTask>,
+    in_progress: HashMap<String, InProgressTask>,
+    is_tty: bool,
+    terminal_width: usize,
+    max_example_name_len: usize,
+    status_lines_displayed: usize,
+    total_examples: usize,
+}
+
+#[derive(Clone)]
+struct InProgressTask {
+    step: Step,
+    started_at: Instant,
+    substatus: Option<String>,
+    info: Option<(String, InfoStyle)>,
+}
+
+struct CompletedTask {
+    step: Step,
+    example_name: String,
+    duration: Duration,
+    info: Option<(String, InfoStyle)>,
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
+pub enum Step {
+    LoadProject,
+    Context,
+    FormatPrompt,
+    Predict,
+    Score,
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub enum InfoStyle {
+    Normal,
+    Warning,
+}
+
+impl Step {
+    pub fn label(&self) -> &'static str {
+        match self {
+            Step::LoadProject => "Load",
+            Step::Context => "Context",
+            Step::FormatPrompt => "Format",
+            Step::Predict => "Predict",
+            Step::Score => "Score",
+        }
+    }
+
+    fn color_code(&self) -> &'static str {
+        match self {
+            Step::LoadProject => "\x1b[33m",
+            Step::Context => "\x1b[35m",
+            Step::FormatPrompt => "\x1b[34m",
+            Step::Predict => "\x1b[32m",
+            Step::Score => "\x1b[31m",
+        }
+    }
+}
+
+const RIGHT_MARGIN: usize = 4;
+
+impl Progress {
+    pub fn new(total_examples: usize) -> Arc<Self> {
+        Arc::new(Self {
+            inner: Mutex::new(ProgressInner {
+                completed: Vec::new(),
+                in_progress: HashMap::new(),
+                is_tty: std::io::stderr().is_terminal(),
+                terminal_width: get_terminal_width(),
+                max_example_name_len: 0,
+                status_lines_displayed: 0,
+                total_examples,
+            }),
+        })
+    }
+
+    pub fn start(self: &Arc<Self>, step: Step, example_name: &str) -> Arc<StepProgress> {
+        {
+            let mut inner = self.inner.lock().unwrap();
+
+            Self::clear_status_lines(&mut inner);
+
+            inner.max_example_name_len = inner.max_example_name_len.max(example_name.len());
+
+            inner.in_progress.insert(
+                example_name.to_string(),
+                InProgressTask {
+                    step,
+                    started_at: Instant::now(),
+                    substatus: None,
+                    info: None,
+                },
+            );
+
+            Self::print_status_lines(&mut inner);
+        }
+
+        Arc::new(StepProgress {
+            progress: self.clone(),
+            step,
+            example_name: example_name.to_string(),
+        })
+    }
+
+    pub fn finish(&self, step: Step, example_name: &str) {
+        let mut inner = self.inner.lock().unwrap();
+
+        let task = inner.in_progress.remove(example_name);
+        if let Some(task) = task {
+            if task.step == step {
+                inner.completed.push(CompletedTask {
+                    step: task.step,
+                    example_name: example_name.to_string(),
+                    duration: task.started_at.elapsed(),
+                    info: task.info,
+                });
+
+                Self::clear_status_lines(&mut inner);
+                Self::print_completed(&inner, inner.completed.last().unwrap());
+                Self::print_status_lines(&mut inner);
+            } else {
+                inner.in_progress.insert(example_name.to_string(), task);
+            }
+        }
+    }
+
+    fn clear_status_lines(inner: &mut ProgressInner) {
+        if inner.is_tty && inner.status_lines_displayed > 0 {
+            // Move up and clear each line we previously displayed
+            for _ in 0..inner.status_lines_displayed {
+                eprint!("\x1b[A\x1b[K");
+            }
+            let _ = std::io::stderr().flush();
+            inner.status_lines_displayed = 0;
+        }
+    }
+
+    fn print_completed(inner: &ProgressInner, task: &CompletedTask) {
+        let duration = format_duration(task.duration);
+        let name_width = inner.max_example_name_len;
+
+        if inner.is_tty {
+            let reset = "\x1b[0m";
+            let bold = "\x1b[1m";
+            let dim = "\x1b[2m";
+
+            let yellow = "\x1b[33m";
+            let info_part = task
+                .info
+                .as_ref()
+                .map(|(s, style)| {
+                    if *style == InfoStyle::Warning {
+                        format!("{yellow}{s}{reset}")
+                    } else {
+                        s.to_string()
+                    }
+                })
+                .unwrap_or_default();
+
+            let prefix = format!(
+                "{bold}{color}{label:>12}{reset} {name:<name_width$} {dim}│{reset} {info_part}",
+                color = task.step.color_code(),
+                label = task.step.label(),
+                name = task.example_name,
+            );
+
+            let duration_with_margin = format!("{duration} ");
+            let padding_needed = inner
+                .terminal_width
+                .saturating_sub(RIGHT_MARGIN)
+                .saturating_sub(duration_with_margin.len())
+                .saturating_sub(strip_ansi_len(&prefix));
+            let padding = " ".repeat(padding_needed);
+
+            eprintln!("{prefix}{padding}{dim}{duration_with_margin}{reset}");
+        } else {
+            let info_part = task
+                .info
+                .as_ref()
+                .map(|(s, _)| format!(" | {}", s))
+                .unwrap_or_default();
+
+            eprintln!(
+                "{label:>12} {name:<name_width$}{info_part} {duration}",
+                label = task.step.label(),
+                name = task.example_name,
+            );
+        }
+    }
+
+    fn print_status_lines(inner: &mut ProgressInner) {
+        if !inner.is_tty || inner.in_progress.is_empty() {
+            inner.status_lines_displayed = 0;
+            return;
+        }
+
+        let reset = "\x1b[0m";
+        let bold = "\x1b[1m";
+        let dim = "\x1b[2m";
+
+        // Build the done/in-progress/total label
+        let done_count = inner.completed.len();
+        let in_progress_count = inner.in_progress.len();
+        let range_label = format!(
+            " {}/{}/{} ",
+            done_count, in_progress_count, inner.total_examples
+        );
+
+        // Print a divider line with range label aligned with timestamps
+        let range_visible_len = range_label.len();
+        let left_divider_len = inner
+            .terminal_width
+            .saturating_sub(RIGHT_MARGIN)
+            .saturating_sub(range_visible_len);
+        let left_divider = "─".repeat(left_divider_len);
+        let right_divider = "─".repeat(RIGHT_MARGIN);
+        eprintln!("{dim}{left_divider}{reset}{range_label}{dim}{right_divider}{reset}");
+
+        let mut tasks: Vec<_> = inner.in_progress.iter().collect();
+        tasks.sort_by_key(|(name, _)| *name);
+
+        let mut lines_printed = 0;
+
+        for (name, task) in tasks.iter() {
+            let elapsed = format_duration(task.started_at.elapsed());
+            let substatus_part = task
+                .substatus
+                .as_ref()
+                .map(|s| truncate_with_ellipsis(s, 30))
+                .unwrap_or_default();
+
+            let step_label = task.step.label();
+            let step_color = task.step.color_code();
+            let name_width = inner.max_example_name_len;
+
+            let prefix = format!(
+                "{bold}{step_color}{step_label:>12}{reset} {name:<name_width$} {dim}│{reset} {substatus_part}",
+                name = name,
+            );
+
+            let duration_with_margin = format!("{elapsed} ");
+            let padding_needed = inner
+                .terminal_width
+                .saturating_sub(RIGHT_MARGIN)
+                .saturating_sub(duration_with_margin.len())
+                .saturating_sub(strip_ansi_len(&prefix));
+            let padding = " ".repeat(padding_needed);
+
+            eprintln!("{prefix}{padding}{dim}{duration_with_margin}{reset}");
+            lines_printed += 1;
+        }
+
+        inner.status_lines_displayed = lines_printed + 1; // +1 for the divider line
+        let _ = std::io::stderr().flush();
+    }
+
+    pub fn clear(&self) {
+        let mut inner = self.inner.lock().unwrap();
+        Self::clear_status_lines(&mut inner);
+    }
+}
+
+pub struct StepProgress {
+    progress: Arc<Progress>,
+    step: Step,
+    example_name: String,
+}
+
+impl StepProgress {
+    pub fn set_substatus(&self, substatus: impl Into<Cow<'static, str>>) {
+        let mut inner = self.progress.inner.lock().unwrap();
+        if let Some(task) = inner.in_progress.get_mut(&self.example_name) {
+            task.substatus = Some(substatus.into().into_owned());
+            Progress::clear_status_lines(&mut inner);
+            Progress::print_status_lines(&mut inner);
+        }
+    }
+
+    pub fn clear_substatus(&self) {
+        let mut inner = self.progress.inner.lock().unwrap();
+        if let Some(task) = inner.in_progress.get_mut(&self.example_name) {
+            task.substatus = None;
+            Progress::clear_status_lines(&mut inner);
+            Progress::print_status_lines(&mut inner);
+        }
+    }
+
+    pub fn set_info(&self, info: impl Into<String>, style: InfoStyle) {
+        let mut inner = self.progress.inner.lock().unwrap();
+        if let Some(task) = inner.in_progress.get_mut(&self.example_name) {
+            task.info = Some((info.into(), style));
+        }
+    }
+}
+
+impl Drop for StepProgress {
+    fn drop(&mut self) {
+        self.progress.finish(self.step, &self.example_name);
+    }
+}
+
+#[cfg(unix)]
+fn get_terminal_width() -> usize {
+    unsafe {
+        let mut winsize: libc::winsize = std::mem::zeroed();
+        if libc::ioctl(libc::STDERR_FILENO, libc::TIOCGWINSZ, &mut winsize) == 0
+            && winsize.ws_col > 0
+        {
+            winsize.ws_col as usize
+        } else {
+            80
+        }
+    }
+}
+
+#[cfg(not(unix))]
+fn get_terminal_width() -> usize {
+    80
+}
+
+fn strip_ansi_len(s: &str) -> usize {
+    let mut len = 0;
+    let mut in_escape = false;
+    for c in s.chars() {
+        if c == '\x1b' {
+            in_escape = true;
+        } else if in_escape {
+            if c == 'm' {
+                in_escape = false;
+            }
+        } else {
+            len += 1;
+        }
+    }
+    len
+}
+
+fn truncate_with_ellipsis(s: &str, max_len: usize) -> String {
+    if s.len() <= max_len {
+        s.to_string()
+    } else {
+        format!("{}…", &s[..max_len.saturating_sub(1)])
+    }
+}
+
+fn format_duration(duration: Duration) -> String {
+    const MINUTE_IN_MILLIS: f32 = 60. * 1000.;
+
+    let millis = duration.as_millis() as f32;
+    if millis < 1000.0 {
+        format!("{}ms", millis)
+    } else if millis < MINUTE_IN_MILLIS {
+        format!("{:.1}s", millis / 1_000.0)
+    } else {
+        format!("{:.1}m", millis / MINUTE_IN_MILLIS)
+    }
+}