use std::sync::Mutex; use regex::{Captures, Regex}; #[cfg(windows)] use windows_sys::Win32::System::Console::{ CONSOLE_SCREEN_BUFFER_INFO, FOREGROUND_BLUE, FOREGROUND_GREEN, FOREGROUND_INTENSITY, FOREGROUND_RED, GetConsoleScreenBufferInfo, GetStdHandle, STD_OUTPUT_HANDLE, SetConsoleTextAttribute, }; static CONSOLE_LOCK: Mutex<()> = Mutex::new(()); #[derive(Clone, Copy, Debug, Eq, PartialEq)] pub enum ConsoleColor { Default, Green, Red, Yellow, Cyan, DarkGray, } pub struct LineBuffer { buffer: String, filter_python: bool, } impl LineBuffer { pub fn new(filter_python: bool) -> Self { Self { buffer: String::new(), filter_python, } } pub fn append(&mut self, bytes: &[u8], mut process_line: F) where F: FnMut(String), { self.buffer.push_str(&String::from_utf8_lossy(bytes)); while let Some(pos) = self.buffer.find('\n') { let mut line = self.buffer[..pos].to_string(); self.buffer.drain(..=pos); if line.ends_with('\r') { line.pop(); } if self.filter_python && !line.contains("[Python] ") { continue; } process_line(line); } } pub fn flush(&mut self, mut process_line: F) where F: FnMut(String), { if self.buffer.is_empty() { return; } let mut line = std::mem::take(&mut self.buffer); if line.ends_with('\r') { line.pop(); } if self.filter_python && !line.contains("[Python] ") { return; } process_line(line); } } pub fn process_stdout_line(line: &str) { if line.contains(" [INFO][Engine] ") { return; } let color = if line.contains("[INFO][Developer]") { ConsoleColor::DarkGray } else if contains_ignore_ascii_case(line, "SUC") { ConsoleColor::Green } else if contains_ignore_ascii_case(line, "ERROR") { ConsoleColor::Red } else if contains_ignore_ascii_case(line, "WARN") { ConsoleColor::Yellow } else if contains_ignore_ascii_case(line, "DEBUG") { ConsoleColor::Cyan } else { ConsoleColor::Default }; print_colored(line, color); } pub fn process_stderr_line(line: &str) { let line = rewrite_python_traceback_path(line); print_colored(&line, ConsoleColor::Red); } pub fn rewrite_python_traceback_path(line: &str) -> String { let re = Regex::new(r#"File "([A-Za-z0-9_\.]+)", line (\d+)"#).unwrap(); re.replace_all(line, |caps: &Captures<'_>| { format!( "File \"{}.py\", line {}", caps[1].replace('.', "/"), &caps[2] ) }) .into_owned() } pub fn print_colored(message: &str, color: ConsoleColor) { let _guard = CONSOLE_LOCK.lock().unwrap(); #[cfg(windows)] unsafe { let handle = GetStdHandle(STD_OUTPUT_HANDLE); if handle.is_null() { println!("{}", message); return; } let mut info: CONSOLE_SCREEN_BUFFER_INFO = std::mem::zeroed(); let has_info = GetConsoleScreenBufferInfo(handle, &mut info) != 0; if color != ConsoleColor::Default { SetConsoleTextAttribute(handle, color_attr(color)); } println!("{}", message); if color != ConsoleColor::Default && has_info { SetConsoleTextAttribute(handle, info.wAttributes); } } #[cfg(not(windows))] { println!("{}", message); } } #[cfg(windows)] fn color_attr(color: ConsoleColor) -> u16 { match color { ConsoleColor::Default => FOREGROUND_RED | FOREGROUND_GREEN | FOREGROUND_BLUE, ConsoleColor::Green => FOREGROUND_GREEN | FOREGROUND_INTENSITY, ConsoleColor::Red => FOREGROUND_RED | FOREGROUND_INTENSITY, ConsoleColor::Yellow => FOREGROUND_RED | FOREGROUND_GREEN | FOREGROUND_INTENSITY, ConsoleColor::Cyan => FOREGROUND_GREEN | FOREGROUND_BLUE | FOREGROUND_INTENSITY, ConsoleColor::DarkGray => FOREGROUND_INTENSITY, } } fn contains_ignore_ascii_case(haystack: &str, needle: &str) -> bool { haystack .as_bytes() .windows(needle.len()) .any(|window| window.eq_ignore_ascii_case(needle.as_bytes())) } #[cfg(test)] mod tests { use super::{LineBuffer, rewrite_python_traceback_path}; #[test] fn process_buffer_append_handles_crlf_half_lines_and_flush() { let mut buffer = LineBuffer::new(true); let mut lines = Vec::new(); buffer.append(b"noise\r\n[Python] first\r\n[Python] sec", |line| { lines.push(line) }); assert_eq!(lines, vec!["[Python] first"]); buffer.append(b"ond\n", |line| lines.push(line)); buffer.flush(|line| lines.push(line)); assert_eq!(lines, vec!["[Python] first", "[Python] second"]); } #[test] fn rewrite_traceback_module_path_to_file_path() { assert_eq!( rewrite_python_traceback_path(r#"Traceback File "a.b", line 3"#), r#"Traceback File "a/b.py", line 3"# ); } }