From 0042dce9498ed0e0d46e3e7cd86b85dcb822c8e1 Mon Sep 17 00:00:00 2001 From: Andrew Gallant Date: Wed, 7 Sep 2016 21:54:28 -0400 Subject: [PATCH] Hack in Windows console coloring. The code has suffered and needs refactoring/commenting. BUT... IT WORKS! --- src/main.rs | 5 +- src/out.rs | 44 +++++++++- src/printer.rs | 185 ++++++++++++++++++++++++++++++++++++++----- src/search.rs | 5 +- src/search_buffer.rs | 3 +- src/sys.rs | 4 +- 6 files changed, 215 insertions(+), 31 deletions(-) diff --git a/src/main.rs b/src/main.rs index 15b250a6a..722232a8a 100644 --- a/src/main.rs +++ b/src/main.rs @@ -35,6 +35,7 @@ use std::thread; use crossbeam::sync::chase_lev::{self, Steal, Stealer}; use grep::Grep; use memmap::{Mmap, Protection}; +use term::Terminal; use walkdir::DirEntry; use args::Args; @@ -198,11 +199,11 @@ impl Worker { let mut printer = self.args.printer(outbuf); self.do_work(&mut printer, work); let outbuf = printer.into_inner(); - if !outbuf.is_empty() { + if !outbuf.get_ref().is_empty() { let mut out = self.out.lock().unwrap(); out.write(&outbuf); } - self.outbuf = Some(outbuf); + self.outbuf = Some(outbuf.into_inner()); } self.match_count } diff --git a/src/out.rs b/src/out.rs index 7f502eebf..9fb31a093 100644 --- a/src/out.rs +++ b/src/out.rs @@ -1,5 +1,11 @@ use std::io::{self, Write}; +use term::{StdoutTerminal, Terminal}; +#[cfg(windows)] +use term::WinConsole; + +use printer::Writer; + /// Out controls the actual output of all search results for a particular file /// to the end user. /// @@ -8,15 +14,32 @@ use std::io::{self, Write}; /// file as a whole. For example, it knows when to print a file separator.) pub struct Out { wtr: io::BufWriter, + term: Option>, printed: bool, file_separator: Option>, } +/// This is like term::stdout, but on Windows always uses WinConsole instead +/// of trying for a TerminfoTerminal. This may be a mistake. +#[cfg(windows)] +fn term_stdout() -> Option> { + WinConsole::new(io::stdout()) + .ok() + .map(|t| Box::new(t) as Box) +} + +#[cfg(not(windows))] +fn term_stdout() -> Option> { + // We never use this crap on *nix. + None +} + impl Out { /// Create a new Out that writes to the wtr given. pub fn new(wtr: W) -> Out { Out { wtr: io::BufWriter::new(wtr), + term: term_stdout(), printed: false, file_separator: None, } @@ -33,14 +56,31 @@ impl Out { /// Write the search results of a single file to the underlying wtr and /// flush wtr. - pub fn write(&mut self, buf: &[u8]) { + pub fn write(&mut self, buf: &Writer>) { if let Some(ref sep) = self.file_separator { if self.printed { let _ = self.wtr.write_all(sep); let _ = self.wtr.write_all(b"\n"); } } - let _ = self.wtr.write_all(buf); + match *buf { + Writer::Colored(ref tt) => { + let _ = self.wtr.write_all(tt.get_ref()); + } + Writer::Windows(ref w) => { + match self.term { + None => { + let _ = self.wtr.write_all(w.get_ref()); + } + Some(ref mut stdout) => { + w.print_stdout(stdout); + } + } + } + Writer::NoColor(ref buf) => { + let _ = self.wtr.write_all(buf); + } + } let _ = self.wtr.flush(); self.printed = true; } diff --git a/src/printer.rs b/src/printer.rs index 6204d5b5e..b82e0fece 100644 --- a/src/printer.rs +++ b/src/printer.rs @@ -3,7 +3,7 @@ use std::path::Path; use std::sync::Arc; use regex::bytes::Regex; -use term::{self, Terminal}; +use term::{self, StdoutTerminal, Terminal}; use term::color::*; use term::terminfo::TermInfo; @@ -115,9 +115,9 @@ impl Printer { } /// Flushes the underlying writer and returns it. - pub fn into_inner(mut self) -> W { + pub fn into_inner(mut self) -> Writer { let _ = self.wtr.flush(); - self.wtr.into_inner() + self.wtr } /// Prints a type definition. @@ -208,7 +208,7 @@ impl Printer { let mut last_written = 0; for (s, e) in re.find_iter(buf) { self.write(&buf[last_written..s]); - let _ = self.wtr.fg(RED); + let _ = self.wtr.fg(BRIGHT_RED); let _ = self.wtr.attr(term::Attr::Bold); self.write(&buf[s..e]); let _ = self.wtr.reset(); @@ -242,7 +242,7 @@ impl Printer { fn write_heading>(&mut self, path: P) { if self.wtr.is_color() { - let _ = self.wtr.fg(GREEN); + let _ = self.wtr.fg(BRIGHT_GREEN); let _ = self.wtr.attr(term::Attr::Bold); } self.write(path.as_ref().to_string_lossy().as_bytes()); @@ -254,7 +254,7 @@ impl Printer { fn line_number(&mut self, n: u64, sep: u8) { if self.wtr.is_color() { - let _ = self.wtr.fg(BLUE); + let _ = self.wtr.fg(BRIGHT_BLUE); let _ = self.wtr.attr(term::Attr::Bold); } self.write(n.to_string().as_bytes()); @@ -278,11 +278,32 @@ impl Printer { } } -enum Writer { +/// Writer corresponds to the final output buffer for search results. All +/// search results are written to a Writer and then a Writer is flushed to +/// stdout only after the full search has completed. +pub enum Writer { Colored(TerminfoTerminal), + Windows(WindowsWriter), NoColor(W), } +pub struct WindowsWriter { + wtr: W, + pos: usize, + colors: Vec, +} + +pub struct WindowsColor { + pos: usize, + opt: WindowsOption, +} + +pub enum WindowsOption { + Foreground(Color), + Background(Color), + Reset, +} + lazy_static! { static ref TERMINFO: Option> = { match term::terminfo::TermInfo::from_env() { @@ -300,7 +321,13 @@ impl Writer { // If we want color, build a TerminfoTerminal and see if the current // environment supports coloring. If not, bail with NoColor. To avoid // losing our writer (ownership), do this the long way. - if !color || TERMINFO.is_none() { + if !color { + return NoColor(wtr); + } + if cfg!(windows) { + return Windows(WindowsWriter { wtr: wtr, pos: 0, colors: vec![] }); + } + if TERMINFO.is_none() { return NoColor(wtr); } let info = TERMINFO.clone().unwrap(); @@ -315,28 +342,35 @@ impl Writer { fn is_color(&self) -> bool { match *self { Colored(_) => true, + Windows(_) => true, NoColor(_) => false, } } - fn map_result( + fn map_result( &mut self, mut f: F, + mut g: G, ) -> term::Result<()> - where F: FnMut(&mut TerminfoTerminal) -> term::Result<()> { + where F: FnMut(&mut TerminfoTerminal) -> term::Result<()>, + G: FnMut(&mut WindowsWriter) -> term::Result<()> { match *self { Colored(ref mut w) => f(w), + Windows(ref mut w) => g(w), NoColor(_) => Err(term::Error::NotSupported), } } - fn map_bool( + fn map_bool( &self, mut f: F, + mut g: G, ) -> bool - where F: FnMut(&TerminfoTerminal) -> bool { + where F: FnMut(&TerminfoTerminal) -> bool, + G: FnMut(&WindowsWriter) -> bool { match *self { Colored(ref w) => f(w), + Windows(ref w) => g(w), NoColor(_) => false, } } @@ -346,6 +380,7 @@ impl io::Write for Writer { fn write(&mut self, buf: &[u8]) -> io::Result { match *self { Colored(ref mut w) => w.write(buf), + Windows(ref mut w) => w.write(buf), NoColor(ref mut w) => w.write(buf), } } @@ -353,6 +388,7 @@ impl io::Write for Writer { fn flush(&mut self) -> io::Result<()> { match *self { Colored(ref mut w) => w.flush(), + Windows(ref mut w) => w.flush(), NoColor(ref mut w) => w.flush(), } } @@ -362,48 +398,49 @@ impl term::Terminal for Writer { type Output = W; fn fg(&mut self, fg: term::color::Color) -> term::Result<()> { - self.map_result(|w| w.fg(fg)) + self.map_result(|w| w.fg(fg), |w| w.fg(fg)) } fn bg(&mut self, bg: term::color::Color) -> term::Result<()> { - self.map_result(|w| w.bg(bg)) + self.map_result(|w| w.bg(bg), |w| w.bg(bg)) } fn attr(&mut self, attr: term::Attr) -> term::Result<()> { - self.map_result(|w| w.attr(attr)) + self.map_result(|w| w.attr(attr), |w| w.attr(attr)) } fn supports_attr(&self, attr: term::Attr) -> bool { - self.map_bool(|w| w.supports_attr(attr)) + self.map_bool(|w| w.supports_attr(attr), |w| w.supports_attr(attr)) } fn reset(&mut self) -> term::Result<()> { - self.map_result(|w| w.reset()) + self.map_result(|w| w.reset(), |w| w.reset()) } fn supports_reset(&self) -> bool { - self.map_bool(|w| w.supports_reset()) + self.map_bool(|w| w.supports_reset(), |w| w.supports_reset()) } fn supports_color(&self) -> bool { - self.map_bool(|w| w.supports_color()) + self.map_bool(|w| w.supports_color(), |w| w.supports_color()) } fn cursor_up(&mut self) -> term::Result<()> { - self.map_result(|w| w.cursor_up()) + self.map_result(|w| w.cursor_up(), |w| w.cursor_up()) } fn delete_line(&mut self) -> term::Result<()> { - self.map_result(|w| w.delete_line()) + self.map_result(|w| w.delete_line(), |w| w.delete_line()) } fn carriage_return(&mut self) -> term::Result<()> { - self.map_result(|w| w.carriage_return()) + self.map_result(|w| w.carriage_return(), |w| w.carriage_return()) } fn get_ref(&self) -> &W { match *self { Colored(ref w) => w.get_ref(), + Windows(ref w) => w.get_ref(), NoColor(ref w) => w, } } @@ -411,6 +448,7 @@ impl term::Terminal for Writer { fn get_mut(&mut self) -> &mut W { match *self { Colored(ref mut w) => w.get_mut(), + Windows(ref mut w) => w.get_mut(), NoColor(ref mut w) => w, } } @@ -418,7 +456,110 @@ impl term::Terminal for Writer { fn into_inner(self) -> W { match self { Colored(w) => w.into_inner(), + Windows(w) => w.into_inner(), NoColor(w) => w, } } } + +impl WindowsWriter { + fn push(&mut self, opt: WindowsOption) { + let pos = self.pos; + self.colors.push(WindowsColor { pos: pos, opt: opt }); + } +} + +impl WindowsWriter> { + /// Print the contents to the given terminal. + pub fn print_stdout(&self, tt: &mut Box) { + let mut last = 0; + for col in &self.colors { + let _ = tt.write_all(&self.wtr[last..col.pos]); + match col.opt { + WindowsOption::Foreground(c) => { + let _ = tt.fg(c); + } + WindowsOption::Background(c) => { + let _ = tt.bg(c); + } + WindowsOption::Reset => { + let _ = tt.reset(); + } + } + last = col.pos; + } + let _ = tt.write_all(&self.wtr[last..]); + let _ = tt.flush(); + } +} + +impl io::Write for WindowsWriter { + fn write(&mut self, buf: &[u8]) -> io::Result { + let n = try!(self.wtr.write(buf)); + self.pos += n; + Ok(n) + } + + fn flush(&mut self) -> io::Result<()> { + self.wtr.flush() + } +} + +impl term::Terminal for WindowsWriter { + type Output = W; + + fn fg(&mut self, fg: term::color::Color) -> term::Result<()> { + self.push(WindowsOption::Foreground(fg)); + Ok(()) + } + + fn bg(&mut self, bg: term::color::Color) -> term::Result<()> { + self.push(WindowsOption::Background(bg)); + Ok(()) + } + + fn attr(&mut self, attr: term::Attr) -> term::Result<()> { + Err(term::Error::NotSupported) + } + + fn supports_attr(&self, attr: term::Attr) -> bool { + false + } + + fn reset(&mut self) -> term::Result<()> { + self.push(WindowsOption::Reset); + Ok(()) + } + + fn supports_reset(&self) -> bool { + true + } + + fn supports_color(&self) -> bool { + true + } + + fn cursor_up(&mut self) -> term::Result<()> { + Err(term::Error::NotSupported) + } + + fn delete_line(&mut self) -> term::Result<()> { + Err(term::Error::NotSupported) + } + + fn carriage_return(&mut self) -> term::Result<()> { + Err(term::Error::NotSupported) + } + + fn get_ref(&self) -> &W { + &self.wtr + } + + fn get_mut(&mut self) -> &mut W { + &mut self.wtr + } + + fn into_inner(self) -> W { + self.wtr + } +} diff --git a/src/search.rs b/src/search.rs index 4f8ae10d6..cb8eac5e1 100644 --- a/src/search.rs +++ b/src/search.rs @@ -687,6 +687,7 @@ mod tests { use std::path::Path; use grep::{Grep, GrepBuilder}; + use term::Terminal; use printer::Printer; @@ -745,7 +746,7 @@ fn main() { &mut inp, &mut pp, &grep, test_path(), hay(haystack)); map(searcher).run().unwrap() }; - (count, String::from_utf8(pp.into_inner()).unwrap()) + (count, String::from_utf8(pp.into_inner().into_inner()).unwrap()) } fn search TestSearcher>( @@ -761,7 +762,7 @@ fn main() { &mut inp, &mut pp, &grep, test_path(), hay(haystack)); map(searcher).run().unwrap() }; - (count, String::from_utf8(pp.into_inner()).unwrap()) + (count, String::from_utf8(pp.into_inner().into_inner()).unwrap()) } #[test] diff --git a/src/search_buffer.rs b/src/search_buffer.rs index 48ff1ba04..24ece4474 100644 --- a/src/search_buffer.rs +++ b/src/search_buffer.rs @@ -144,6 +144,7 @@ mod tests { use std::path::Path; use grep::{Grep, GrepBuilder}; + use term::Terminal; use printer::Printer; @@ -197,7 +198,7 @@ fn main() { &mut pp, &grep, test_path(), haystack.as_bytes()); map(searcher).run() }; - (count, String::from_utf8(pp.into_inner()).unwrap()) + (count, String::from_utf8(pp.into_inner().into_inner()).unwrap()) } #[test] diff --git a/src/sys.rs b/src/sys.rs index 15205cc73..990097925 100644 --- a/src/sys.rs +++ b/src/sys.rs @@ -5,15 +5,15 @@ redirected to a file? etc... We use this information to tweak various default configuration parameters such as colors and match formatting. */ -use libc; - #[cfg(unix)] pub fn stdin_is_atty() -> bool { + use libc; 0 < unsafe { libc::isatty(libc::STDIN_FILENO) } } #[cfg(unix)] pub fn stdout_is_atty() -> bool { + use libc; 0 < unsafe { libc::isatty(libc::STDOUT_FILENO) } }