Merge pull request #2137 from clippered/add-color-cli-option

add cli option for color

src/bin/rustfmt.rs

@@ -25,7 +25,7 @@ use getopts::{HasArg, Matches, Occur, Options};
 
 use rustfmt::{run, Input, Summary};
 use rustfmt::file_lines::FileLines;
-use rustfmt::config::{get_toml_path, Config, WriteMode};
+use rustfmt::config::{get_toml_path, Color, Config, WriteMode};
 
 type FmtError = Box<error::Error + Send + Sync>;
 type FmtResult<T> = std::result::Result<T, FmtError>;
@@ -59,6 +59,7 @@ struct CliOptions {
     skip_children: bool,
     verbose: bool,
     write_mode: Option<WriteMode>,
+    color: Option<Color>,
     file_lines: FileLines, // Default is all lines in all files.
     unstable_features: bool,
 }
@@ -90,6 +91,13 @@ impl CliOptions {
             }
         }
 
+        if let Some(ref color) = matches.opt_str("color") {
+            match Color::from_str(color) {
+                Ok(color) => options.color = Some(color),
+                _ => return Err(FmtError::from(format!("Invalid color: {}", color))),
+            }
+        }
+
         if let Some(ref file_lines) = matches.opt_str("file-lines") {
             options.file_lines = file_lines.parse()?;
         }
@@ -105,6 +113,9 @@ impl CliOptions {
         if let Some(write_mode) = self.write_mode {
             config.set().write_mode(write_mode);
         }
+        if let Some(color) = self.color {
+            config.set().color(color);
+        }
     }
 }
 
@@ -131,6 +142,12 @@ fn make_opts() -> Options {
         "how to write output (not usable when piping from stdin)",
         "[replace|overwrite|display|plain|diff|coverage|checkstyle]",
     );
+    opts.optopt(
+        "",
+        "color",
+        "use colored output (if supported)",
+        "[always|never|auto]",
+    );
     opts.optflag("", "skip-children", "don't reformat child modules");
 
     opts.optflag(

src/config.rs

@@ -147,6 +147,15 @@ configuration_option_enum! { WriteMode:
     Checkstyle,
 }
 
+configuration_option_enum! { Color:
+    // Always use color, whether it is a piped or terminal output
+    Always,
+    // Never use color
+    Never,
+    // Automatically use color, if supported by terminal
+    Auto,
+}
+
 /// Trait for types that can be used in `Config`.
 pub trait ConfigType: Sized {
     /// Returns hint text for use in `Config::print_docs()`. For enum types, this is a
@@ -628,6 +637,8 @@ create_config! {
     write_mode: WriteMode, WriteMode::Overwrite, false,
         "What Write Mode to use when none is supplied: \
          Replace, Overwrite, Display, Plain, Diff, Coverage";
+    color: Color, Color::Auto, false,
+        "What Color option to use when none is supplied: Always, Never, Auto";
     condense_wildcard_suffixes: bool, false, false, "Replace strings of _ wildcards by a single .. \
                                               in tuple patterns";
     combine_control_expr: bool, true, false, "Combine control expressions with function calls.";

src/filemap.rs

@@ -152,9 +152,11 @@ where
             if let Ok((ori, fmt)) = source_and_formatted_text(text, filename, config) {
                 let mismatch = make_diff(&ori, &fmt, 3);
                 let has_diff = !mismatch.is_empty();
-                print_diff(mismatch, |line_num| {
-                    format!("Diff in {} at line {}:", filename, line_num)
-                });
+                print_diff(
+                    mismatch,
+                    |line_num| format!("Diff in {} at line {}:", filename, line_num),
+                    config.color(),
+                );
                 return Ok(has_diff);
             }
         }

src/lib.rs

@@ -44,7 +44,7 @@ use checkstyle::{output_footer, output_header};
 use config::Config;
 use filemap::FileMap;
 use issues::{BadIssueSeeker, Issue};
-use utils::isatty;
+use utils::use_colored_tty;
 use visitor::FmtVisitor;
 
 pub use self::summary::Summary;
@@ -581,7 +581,8 @@ pub fn run(input: Input, config: &Config) -> Summary {
             if report.has_warnings() {
                 match term::stderr() {
                     Some(ref t)
-                        if isatty() && t.supports_color() && t.supports_attr(term::Attr::Bold) =>
+                        if use_colored_tty(config.color()) && t.supports_color()
+                            && t.supports_attr(term::Attr::Bold) =>
                     {
                         match report.print_warnings_fancy(term::stderr().unwrap()) {
                             Ok(..) => (),

src/rustfmt_diff.rs

@@ -8,11 +8,12 @@
 // option. This file may not be copied, modified, or distributed
 // except according to those terms.
 
+use config::Color;
 use diff;
 use std::collections::VecDeque;
 use std::io;
 use term;
-use utils::isatty;
+use utils::use_colored_tty;
 
 #[derive(Debug, PartialEq)]
 pub enum DiffLine {
@@ -96,12 +97,12 @@ pub fn make_diff(expected: &str, actual: &str, context_size: usize) -> Vec<Misma
     results
 }
 
-pub fn print_diff<F>(diff: Vec<Mismatch>, get_section_title: F)
+pub fn print_diff<F>(diff: Vec<Mismatch>, get_section_title: F, color: Color)
 where
     F: Fn(u32) -> String,
 {
     match term::stdout() {
-        Some(ref t) if isatty() && t.supports_color() => {
+        Some(ref t) if use_colored_tty(color) && t.supports_color() => {
             print_diff_fancy(diff, get_section_title, term::stdout().unwrap())
         }
         _ => print_diff_basic(diff, get_section_title),

src/utils.rs

@@ -15,6 +15,7 @@ use syntax::ast::{self, Attribute, CrateSugar, MetaItem, MetaItemKind, NestedMet
                   NestedMetaItemKind, Path, Visibility};
 use syntax::codemap::{BytePos, Span, NO_EXPANSION};
 
+use config::Color;
 use rewrite::RewriteContext;
 use shape::Shape;
 
@@ -484,6 +485,14 @@ pub fn isatty() -> bool {
     }
 }
 
+pub fn use_colored_tty(color: Color) -> bool {
+    match color {
+        Color::Always => true,
+        Color::Never => false,
+        Color::Auto => isatty(),
+    }
+}
+
 pub fn starts_with_newline(s: &str) -> bool {
     s.starts_with('\n') || s.starts_with("\r\n")
 }

tests/system.rs

@@ -23,7 +23,7 @@ use std::str::Chars;
 
 use rustfmt::*;
 use rustfmt::filemap::{write_system_newlines, FileMap};
-use rustfmt::config::{Config, ReportTactic};
+use rustfmt::config::{Color, Config, ReportTactic};
 use rustfmt::rustfmt_diff::*;
 
 const DIFF_CONTEXT_SIZE: usize = 3;
@@ -229,9 +229,11 @@ fn print_mismatches(result: HashMap<String, Vec<Mismatch>>) {
     let mut t = term::stdout().unwrap();
 
     for (file_name, diff) in result {
-        print_diff(diff, |line_num| {
-            format!("\nMismatch at {}:{}:", file_name, line_num)
-        });
+        print_diff(
+            diff,
+            |line_num| format!("\nMismatch at {}:{}:", file_name, line_num),
+            Color::Auto,
+        );
     }
 
     t.reset().unwrap();