Hack in Windows console coloring.

The code has suffered and needs refactoring/commenting. BUT... IT WORKS!
This commit is contained in:
Andrew Gallant 2016-09-07 21:54:28 -04:00
parent ca058d7584
commit 0042dce949
6 changed files with 215 additions and 31 deletions

View File

@ -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
}

View File

@ -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<W: io::Write> {
wtr: io::BufWriter<W>,
term: Option<Box<StdoutTerminal>>,
printed: bool,
file_separator: Option<Vec<u8>>,
}
/// 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<Box<StdoutTerminal>> {
WinConsole::new(io::stdout())
.ok()
.map(|t| Box::new(t) as Box<StdoutTerminal>)
}
#[cfg(not(windows))]
fn term_stdout() -> Option<Box<StdoutTerminal>> {
// We never use this crap on *nix.
None
}
impl<W: io::Write> Out<W> {
/// Create a new Out that writes to the wtr given.
pub fn new(wtr: W) -> Out<W> {
Out {
wtr: io::BufWriter::new(wtr),
term: term_stdout(),
printed: false,
file_separator: None,
}
@ -33,14 +56,31 @@ impl<W: io::Write> Out<W> {
/// 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<Vec<u8>>) {
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;
}

View File

@ -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<W: Send + io::Write> Printer<W> {
}
/// Flushes the underlying writer and returns it.
pub fn into_inner(mut self) -> W {
pub fn into_inner(mut self) -> Writer<W> {
let _ = self.wtr.flush();
self.wtr.into_inner()
self.wtr
}
/// Prints a type definition.
@ -208,7 +208,7 @@ impl<W: Send + io::Write> Printer<W> {
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<W: Send + io::Write> Printer<W> {
fn write_heading<P: AsRef<Path>>(&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<W: Send + io::Write> Printer<W> {
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<W: Send + io::Write> Printer<W> {
}
}
enum Writer<W> {
/// 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<W> {
Colored(TerminfoTerminal<W>),
Windows(WindowsWriter<W>),
NoColor(W),
}
pub struct WindowsWriter<W> {
wtr: W,
pos: usize,
colors: Vec<WindowsColor>,
}
pub struct WindowsColor {
pos: usize,
opt: WindowsOption,
}
pub enum WindowsOption {
Foreground(Color),
Background(Color),
Reset,
}
lazy_static! {
static ref TERMINFO: Option<Arc<TermInfo>> = {
match term::terminfo::TermInfo::from_env() {
@ -300,7 +321,13 @@ impl<W: Send + io::Write> Writer<W> {
// 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<W: Send + io::Write> Writer<W> {
fn is_color(&self) -> bool {
match *self {
Colored(_) => true,
Windows(_) => true,
NoColor(_) => false,
}
}
fn map_result<F>(
fn map_result<F, G>(
&mut self,
mut f: F,
mut g: G,
) -> term::Result<()>
where F: FnMut(&mut TerminfoTerminal<W>) -> term::Result<()> {
where F: FnMut(&mut TerminfoTerminal<W>) -> term::Result<()>,
G: FnMut(&mut WindowsWriter<W>) -> term::Result<()> {
match *self {
Colored(ref mut w) => f(w),
Windows(ref mut w) => g(w),
NoColor(_) => Err(term::Error::NotSupported),
}
}
fn map_bool<F>(
fn map_bool<F, G>(
&self,
mut f: F,
mut g: G,
) -> bool
where F: FnMut(&TerminfoTerminal<W>) -> bool {
where F: FnMut(&TerminfoTerminal<W>) -> bool,
G: FnMut(&WindowsWriter<W>) -> bool {
match *self {
Colored(ref w) => f(w),
Windows(ref w) => g(w),
NoColor(_) => false,
}
}
@ -346,6 +380,7 @@ impl<W: Send + io::Write> io::Write for Writer<W> {
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
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<W: Send + io::Write> io::Write for Writer<W> {
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<W: Send + io::Write> term::Terminal for Writer<W> {
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<W: Send + io::Write> term::Terminal for Writer<W> {
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<W: Send + io::Write> term::Terminal for Writer<W> {
fn into_inner(self) -> W {
match self {
Colored(w) => w.into_inner(),
Windows(w) => w.into_inner(),
NoColor(w) => w,
}
}
}
impl<W: Send + io::Write> WindowsWriter<W> {
fn push(&mut self, opt: WindowsOption) {
let pos = self.pos;
self.colors.push(WindowsColor { pos: pos, opt: opt });
}
}
impl WindowsWriter<Vec<u8>> {
/// Print the contents to the given terminal.
pub fn print_stdout(&self, tt: &mut Box<StdoutTerminal>) {
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<W: Send + io::Write> io::Write for WindowsWriter<W> {
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
let n = try!(self.wtr.write(buf));
self.pos += n;
Ok(n)
}
fn flush(&mut self) -> io::Result<()> {
self.wtr.flush()
}
}
impl<W: Send + io::Write> term::Terminal for WindowsWriter<W> {
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
}
}

View File

@ -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<F: FnMut(TestSearcher) -> 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]

View File

@ -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]

View File

@ -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) }
}