use std::io::Write;
pub(super) struct LinePrefixer {
label: String,
pending: Vec<u8>,
}
impl LinePrefixer {
pub(super) fn new(label: &str, prefix: bool, allow_color: bool) -> Self {
if !prefix {
return Self {
label: String::new(),
pending: Vec::new(),
};
}
let plain = format!("{label} | ");
let label = crate::ui::paint(
crate::ui::service_style(label),
&plain,
allow_color && crate::ui::stdout_colored(),
);
Self {
label,
pending: Vec::new(),
}
}
pub(super) fn write(&mut self, out: &mut impl Write, chunk: &[u8]) {
self.pending.extend_from_slice(chunk);
while let Some(nl) = self.pending.iter().position(|&b| b == b'\n') {
let _ = out.write_all(self.label.as_bytes());
let _ = out.write_all(&self.pending[..=nl]);
self.pending.drain(..=nl);
}
let _ = out.flush();
}
pub(super) fn flush_tail(&mut self, out: &mut impl Write) {
if !self.pending.is_empty() {
let _ = out.write_all(self.label.as_bytes());
let _ = out.write_all(&self.pending);
let _ = out.write_all(b"\n");
let _ = out.flush();
self.pending.clear();
}
}
}
#[cfg(test)]
mod tests {
use super::LinePrefixer;
#[test]
fn line_prefixer_tags_lines_and_buffers_partials() {
let mut p = LinePrefixer::new("web", true, false);
let mut out: Vec<u8> = Vec::new();
p.write(&mut out, b"hello\nwor");
assert_eq!(out, b"web | hello\n");
p.write(&mut out, b"ld\n");
assert_eq!(out, b"web | hello\nweb | world\n");
}
#[test]
fn line_prefixer_flush_tail_emits_unterminated_line() {
let mut p = LinePrefixer::new("db", true, false);
let mut out: Vec<u8> = Vec::new();
p.write(&mut out, b"partial");
assert!(out.is_empty(), "a line with no newline is held back");
p.flush_tail(&mut out);
assert_eq!(out, b"db | partial\n");
}
#[test]
fn line_prefixer_no_prefix_emits_bare_lines() {
let mut p = LinePrefixer::new("web", false, false);
let mut out: Vec<u8> = Vec::new();
p.write(&mut out, b"hello\n");
assert_eq!(out, b"hello\n");
p.write(&mut out, b"tail");
p.flush_tail(&mut out);
assert_eq!(out, b"hello\ntail\n");
}
}