use std::cell::RefCell;
use std::fmt;
use std::fs::OpenOptions;
use std::io::{self, Write};
use std::path::Path;
use std::sync::{LazyLock, Mutex};
thread_local! {
static CAPTURE: RefCell<Option<Vec<u8>>> = const { RefCell::new(None) };
}
static LOG_SINK: LazyLock<Mutex<Option<std::fs::File>>> = LazyLock::new(|| Mutex::new(None));
pub fn open_log(path: impl AsRef<Path>, append: bool) -> io::Result<()> {
let file = OpenOptions::new()
.create(true)
.write(true)
.truncate(!append)
.append(append)
.open(path)?;
let mut sink = LOG_SINK
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner());
*sink = Some(file);
Ok(())
}
pub fn close_log() {
let mut sink = LOG_SINK
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner());
if let Some(mut file) = sink.take() {
let _ = file.flush();
}
}
fn log_text(text: &str) {
let mut sink = LOG_SINK
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner());
let Some(file) = sink.as_mut() else {
return;
};
let _ = file.write_all(text.as_bytes());
let _ = file.flush();
}
fn log_args(args: fmt::Arguments<'_>) {
let mut sink = LOG_SINK
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner());
let Some(file) = sink.as_mut() else {
return;
};
let mut text = String::new();
let _ = fmt::write(&mut text, args);
let clean = strip_ansi(&text);
let _ = file.write_all(clean.as_bytes());
let _ = file.flush();
}
pub fn log_input_line(line: &str) {
let mut text = String::with_capacity(line.len() + 3);
text.push_str("> ");
text.push_str(line);
text.push('\n');
let clean = strip_ansi(&text);
log_text(&clean);
}
pub fn write_fmt(args: fmt::Arguments<'_>) {
log_args(args);
let captured = CAPTURE.with(|slot| {
let mut slot = slot.borrow_mut();
let Some(buf) = slot.as_mut() else {
return false;
};
let _ = buf.write_fmt(args);
true
});
if !captured {
print!("{args}");
}
}
pub fn write_stderr_fmt(args: fmt::Arguments<'_>) {
log_args(args);
eprint!("{args}");
}
pub fn capturing() -> bool {
CAPTURE.with(|slot| slot.borrow().is_some())
}
macro_rules! out {
($($arg:tt)*) => {
$crate::output::write_fmt(format_args!($($arg)*))
};
}
macro_rules! outln {
() => {
$crate::output::write_fmt(format_args!("\n"))
};
($($arg:tt)*) => {
$crate::output::write_fmt(format_args!("{}\n", format_args!($($arg)*)))
};
}
struct Restore(Option<Vec<u8>>);
impl Drop for Restore {
fn drop(&mut self) {
CAPTURE.with(|slot| *slot.borrow_mut() = self.0.take());
}
}
pub fn capture<R>(f: impl FnOnce() -> R) -> (R, String) {
let previous = CAPTURE.with(|slot| slot.borrow_mut().replace(Vec::new()));
let restore = Restore(previous);
let result = f();
let buf = CAPTURE
.with(|slot| slot.borrow_mut().take())
.unwrap_or_default();
drop(restore);
(result, strip_ansi(&String::from_utf8_lossy(&buf)))
}
fn strip_ansi(text: &str) -> String {
let mut out = String::with_capacity(text.len());
let mut chars = text.chars().peekable();
while let Some(c) = chars.next() {
if c != '\x1b' {
out.push(c);
continue;
}
match chars.next() {
Some('[') => {
for c in chars.by_ref() {
if ('\x40'..='\x7e').contains(&c) {
break;
}
}
}
Some(']') => {
let mut prev = '\0';
for c in chars.by_ref() {
if c == '\x07' || (prev == '\x1b' && c == '\\') {
break;
}
prev = c;
}
}
Some(other) => out.push(other),
None => {}
}
}
out
}
#[cfg(test)]
mod tests {
use super::*;
use owo_colors::OwoColorize;
#[test]
fn capture_collects_and_strips_styling() {
let (value, text) = capture(|| {
outln!("{} {}", "bold".bold(), "red".bright_red());
out!("tail");
7
});
assert_eq!(value, 7);
assert_eq!(text, "bold red\ntail");
}
#[test]
fn captures_nest_and_restore() {
let (_, outer) = capture(|| {
outln!("before");
let (_, inner) = capture(|| outln!("inner"));
assert_eq!(inner, "inner\n");
outln!("after");
});
assert_eq!(outer, "before\nafter\n");
}
#[test]
fn strip_ansi_handles_osc_and_lone_escape() {
assert_eq!(strip_ansi("a\x1b]0;title\x07b"), "ab");
assert_eq!(strip_ansi("a\x1b]0;title\x1b\\b"), "ab");
assert_eq!(strip_ansi("a\x1b"), "a");
assert_eq!(strip_ansi("plain"), "plain");
}
}