add cli option for color

6 files changed, 51 insertions, 9 deletions

diff --git a/src/bin/rustfmt.rs b/src/bin/rustfmt.rs
index f0e4f0a7897..fd9f9021484 100644
--- a/src/bin/rustfmt.rs
+++ b/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,14 @@ impl CliOptions {
             }
         }
 
+        if let Some(ref color) = matches.opt_str("color") {
+            if let Ok(color) = Color::from_str(color) {
+                options.color = Some(color);
+            } else {
+                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 +114,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 +143,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(
diff --git a/src/config.rs b/src/config.rs
index 98939e9e191..430402d2784 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -152,6 +152,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
@@ -643,6 +652,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.";
diff --git a/src/filemap.rs b/src/filemap.rs
index 72724da1c11..5942934d877 100644
--- a/src/filemap.rs
+++ b/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);
             }
         }
diff --git a/src/lib.rs b/src/lib.rs
index 771fe785339..dfdf4b5143f 100644
--- a/src/lib.rs
+++ b/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::iscolored;
 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 iscolored(config.color()) && t.supports_color()
+                            && t.supports_attr(term::Attr::Bold) =>
                     {
                         match report.print_warnings_fancy(term::stderr().unwrap()) {
                             Ok(..) => (),
diff --git a/src/rustfmt_diff.rs b/src/rustfmt_diff.rs
index c0aca68fe2f..858fb876f6e 100644
--- a/src/rustfmt_diff.rs
+++ b/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::iscolored;
 
 #[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 iscolored(color) && t.supports_color() => {
             print_diff_fancy(diff, get_section_title, term::stdout().unwrap())
         }
         _ => print_diff_basic(diff, get_section_title),
diff --git a/src/utils.rs b/src/utils.rs
index 999ba9f67b9..68fca8789d3 100644
--- a/src/utils.rs
+++ b/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 iscolored(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")
 }