126 lines
3.8 KiB
Rust
126 lines
3.8 KiB
Rust
//! Shared helpers for capturing subprocess output line by line
|
|
//!
|
|
//! Used by the context drivers implementing
|
|
//! [`ContextDriver::run_captured`](super::api::ContextDriver::run_captured):
|
|
//! raw bytes are read incrementally, split into lines, cleaned up (ANSI escape
|
|
//! stripping, carriage-return progress collapsing) and forwarded to a
|
|
//! [`LineSink`](super::api::LineSink).
|
|
|
|
use std::io::Read;
|
|
use std::sync::OnceLock;
|
|
|
|
use super::api::{LineSink, Stream};
|
|
use regex::Regex;
|
|
|
|
/// Precompiled regex matching ANSI escape sequences (CSI and simple escapes)
|
|
fn ansi_regex() -> &'static Regex {
|
|
static RE: OnceLock<Regex> = OnceLock::new();
|
|
RE.get_or_init(|| Regex::new(r"\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])").unwrap())
|
|
}
|
|
|
|
/// Strip ANSI escape sequences from a line
|
|
pub(crate) fn strip_ansi(line: &str) -> String {
|
|
ansi_regex().replace_all(line, "").to_string()
|
|
}
|
|
|
|
/// Read `reader` to EOF, forwarding each decoded line to `sink`
|
|
///
|
|
/// Lines are split on `\n`; when a line contains carriage returns (e.g. apt's
|
|
/// `0% [Working]` progress fragments), only the last `\r`-segment is kept so
|
|
/// progress updates replace each other instead of accumulating. Empty lines
|
|
/// are dropped.
|
|
pub(crate) fn pump<R: Read>(mut reader: R, stream: Stream, sink: &dyn LineSink) {
|
|
let mut buf: Vec<u8> = Vec::with_capacity(8192);
|
|
let mut chunk = [0u8; 4096];
|
|
|
|
loop {
|
|
match reader.read(&mut chunk) {
|
|
Ok(0) => break,
|
|
Ok(n) => buf.extend_from_slice(&chunk[..n]),
|
|
Err(_) => break,
|
|
}
|
|
|
|
while let Some(pos) = buf.iter().position(|&b| b == b'\n') {
|
|
let line: Vec<u8> = buf.drain(..=pos).collect();
|
|
emit(&line[..line.len() - 1], stream, sink);
|
|
}
|
|
}
|
|
|
|
// Flush a trailing line without newline, if any
|
|
if !buf.is_empty() {
|
|
emit(&buf, stream, sink);
|
|
}
|
|
}
|
|
|
|
/// Clean up and forward one raw line to the sink
|
|
fn emit(raw: &[u8], stream: Stream, sink: &dyn LineSink) {
|
|
let mut line = String::from_utf8_lossy(raw).to_string();
|
|
|
|
// Strip a trailing carriage return left over from CRLF line endings
|
|
if line.ends_with('\r') {
|
|
line.pop();
|
|
}
|
|
|
|
// Carriage-return progress: keep only the last segment of the line
|
|
if let Some(idx) = line.rfind('\r') {
|
|
line = line[idx + 1..].to_string();
|
|
}
|
|
|
|
let line = strip_ansi(&line);
|
|
let trimmed = line.trim_end();
|
|
if trimmed.trim().is_empty() {
|
|
return;
|
|
}
|
|
|
|
sink.line(stream, trimmed);
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use std::sync::Mutex;
|
|
|
|
#[derive(Default)]
|
|
struct CollectingSink(Mutex<Vec<(Stream, String)>>);
|
|
|
|
impl LineSink for CollectingSink {
|
|
fn line(&self, stream: Stream, line: &str) {
|
|
self.0.lock().unwrap().push((stream, line.to_string()));
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_pump_splits_lines_and_merges_cr_fragments() {
|
|
let sink = CollectingSink::default();
|
|
let data =
|
|
b"Get:1 http://x InRelease [1 kB]\r0% [Working]\r\nHit:2 http://y Release\npartial";
|
|
pump(&data[..], Stream::Stdout, &sink);
|
|
|
|
let lines = sink.0.lock().unwrap().clone();
|
|
assert_eq!(
|
|
lines,
|
|
vec![
|
|
(Stream::Stdout, "0% [Working]".to_string()),
|
|
(Stream::Stdout, "Hit:2 http://y Release".to_string()),
|
|
(Stream::Stdout, "partial".to_string()),
|
|
]
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_pump_strips_ansi_and_skips_empty_lines() {
|
|
let sink = CollectingSink::default();
|
|
let data = b"\x1b[1mSetting up foo\x1b[0m\n\n \nE: boom\n";
|
|
pump(&data[..], Stream::Stderr, &sink);
|
|
|
|
let lines = sink.0.lock().unwrap().clone();
|
|
assert_eq!(
|
|
lines,
|
|
vec![
|
|
(Stream::Stderr, "Setting up foo".to_string()),
|
|
(Stream::Stderr, "E: boom".to_string()),
|
|
]
|
|
);
|
|
}
|
|
}
|