1 files changed, 42 insertions, 61 deletions

diff --git a/src/rustfmt_diff.rs b/src/rustfmt_diff.rs
index f37ab5c16ef..d563e6eee3c 100644
--- a/src/rustfmt_diff.rs
+++ b/src/rustfmt_diff.rs
@@ -37,6 +37,43 @@ impl Mismatch {
     }
 }
 
+// This struct handles writing output to stdout and abstracts away the logic
+// of printing in color, if it's possible in the executing environment.
+pub struct OutputWriter {
+    terminal: Option<Box<term::Terminal<Output = io::Stdout>>>,
+}
+
+impl OutputWriter {
+    // Create a new OutputWriter instance based on the caller's preference
+    // for colorized output and the capabilities of the terminal.
+    pub fn new(color: Color) -> Self {
+        if let Some(t) = term::stdout() {
+            if use_colored_tty(color) && t.supports_color() {
+                return OutputWriter { terminal: Some(t) };
+            }
+        }
+        OutputWriter { terminal: None }
+    }
+
+    // Write output in the optionally specified color. The output is written
+    // in the specified color if this OutputWriter instance contains a
+    // Terminal in its `terminal` field.
+    pub fn writeln(&mut self, msg: &str, color: Option<term::color::Color>) {
+        match &mut self.terminal {
+            Some(ref mut t) => {
+                if let Some(color) = color {
+                    t.fg(color).unwrap();
+                }
+                writeln!(t, "{}", msg).unwrap();
+                if color.is_some() {
+                    t.reset().unwrap();
+                }
+            }
+            None => println!("{}", msg),
+        }
+    }
+}
+
 // Produces a diff between the expected output and actual output of rustfmt.
 pub fn make_diff(expected: &str, actual: &str, context_size: usize) -> Vec<Mismatch> {
     let mut line_number = 1;
@@ -97,80 +134,24 @@ pub fn make_diff(expected: &str, actual: &str, context_size: usize) -> Vec<Misma
     results
 }
 
-// A representation of how to write output.
-pub enum PrintType {
-    Fancy, // want to output color and the terminal supports it
-    Basic, // do not want to output color or the terminal does not support color
-}
-
-impl PrintType {
-    pub fn get(color: Color) -> Self {
-        match term::stdout() {
-            Some(ref t) if use_colored_tty(color) && t.supports_color() => PrintType::Fancy,
-            _ => PrintType::Basic,
-        }
-    }
-}
-
 pub fn print_diff<F>(diff: Vec<Mismatch>, get_section_title: F, color: Color)
 where
     F: Fn(u32) -> String,
 {
-    match PrintType::get(color) {
-        PrintType::Fancy => print_diff_fancy(diff, get_section_title, term::stdout().unwrap()),
-        PrintType::Basic => print_diff_basic(diff, get_section_title),
-    }
-}
-
-fn print_diff_fancy<F>(
-    diff: Vec<Mismatch>,
-    get_section_title: F,
-    mut t: Box<term::Terminal<Output = io::Stdout>>,
-) where
-    F: Fn(u32) -> String,
-{
-    for mismatch in diff {
-        let title = get_section_title(mismatch.line_number);
-        writeln!(t, "{}", title).unwrap();
-
-        for line in mismatch.lines {
-            match line {
-                DiffLine::Context(ref str) => {
-                    t.reset().unwrap();
-                    writeln!(t, " {}⏎", str).unwrap();
-                }
-                DiffLine::Expected(ref str) => {
-                    t.fg(term::color::GREEN).unwrap();
-                    writeln!(t, "+{}⏎", str).unwrap();
-                }
-                DiffLine::Resulting(ref str) => {
-                    t.fg(term::color::RED).unwrap();
-                    writeln!(t, "-{}⏎", str).unwrap();
-                }
-            }
-        }
-        t.reset().unwrap();
-    }
-}
+    let mut writer = OutputWriter::new(color);
 
-pub fn print_diff_basic<F>(diff: Vec<Mismatch>, get_section_title: F)
-where
-    F: Fn(u32) -> String,
-{
     for mismatch in diff {
         let title = get_section_title(mismatch.line_number);
-        println!("{}", title);
+        writer.writeln(&format!("{}", title), None);
 
         for line in mismatch.lines {
             match line {
-                DiffLine::Context(ref str) => {
-                    println!(" {}⏎", str);
-                }
+                DiffLine::Context(ref str) => writer.writeln(&format!(" {}⏎", str), None),
                 DiffLine::Expected(ref str) => {
-                    println!("+{}⏎", str);
+                    writer.writeln(&format!("+{}⏎", str), Some(term::color::GREEN))
                 }
                 DiffLine::Resulting(ref str) => {
-                    println!("-{}⏎", str);
+                    writer.writeln(&format!("-{}⏎", str), Some(term::color::RED))
                 }
             }
         }