From 46a8afe144e0044f11000a94550e8833557ec316 Mon Sep 17 00:00:00 2001
From: Martin von Zweigbergk <martinvonz@google.com>
Date: Wed, 1 Nov 2023 17:53:13 -0700
Subject: [PATCH] formatter: reset colors on early drop

If we create a `ColorFormatter`, add some labels to it, print
something using the configured style, and then return early because of
an error, we would leave the terminal in a bad state. I think many
shells reset color codes after a command returns, but let's still do
our best.
---
 cli/src/formatter.rs        | 45 +++++++++++++++++++++++++++++++++++--
 cli/src/template_builder.rs |  1 +
 cli/src/text_util.rs        |  1 +
 3 files changed, 45 insertions(+), 2 deletions(-)

diff --git a/cli/src/formatter.rs b/cli/src/formatter.rs
index e44ee6151..a215f87fa 100644
--- a/cli/src/formatter.rs
+++ b/cli/src/formatter.rs
@@ -216,7 +216,7 @@ impl Style {
 }
 
 #[derive(Clone, Debug)]
-pub struct ColorFormatter<W> {
+pub struct ColorFormatter<W: Write> {
     output: W,
     rules: Arc<Rules>,
     /// The stack of currently applied labels. These determine the desired
@@ -460,6 +460,15 @@ impl<W: Write> Formatter for ColorFormatter<W> {
     }
 }
 
+impl<W: Write> Drop for ColorFormatter<W> {
+    fn drop(&mut self) {
+        // If a `ColorFormatter` was dropped without popping all labels first (perhaps
+        // because of an error), let's still try to reset any currently active style.
+        self.labels.clear();
+        self.write_new_style().ok();
+    }
+}
+
 /// Like buffered formatter, but records `push`/`pop_label()` calls.
 ///
 /// This allows you to manipulate the recorded data without losing labels.
@@ -646,6 +655,7 @@ mod tests {
             formatter.pop_label().unwrap();
             formatter.write_str("\n").unwrap();
         }
+        drop(formatter);
         insta::assert_snapshot!(String::from_utf8(output).unwrap(), @r###"
          black 
          red 
@@ -682,6 +692,7 @@ mod tests {
         formatter.write_str(" inside ").unwrap();
         formatter.pop_label().unwrap();
         formatter.write_str(" after ").unwrap();
+        drop(formatter);
         insta::assert_snapshot!(String::from_utf8(output).unwrap(), @" before  inside  after ");
     }
 
@@ -726,6 +737,7 @@ mod tests {
         formatter.pop_label().unwrap();
         formatter.pop_label().unwrap();
         formatter.write_str("\n").unwrap();
+        drop(formatter);
         insta::assert_snapshot!(String::from_utf8(output).unwrap(), @r###"
          fg only 
          bg only 
@@ -754,6 +766,7 @@ mod tests {
         formatter.pop_label().unwrap();
         formatter.write_str(" not bold again ").unwrap();
         formatter.pop_label().unwrap();
+        drop(formatter);
         insta::assert_snapshot!(String::from_utf8(output).unwrap(), @" not bold  bold  not bold again ");
     }
 
@@ -776,6 +789,7 @@ mod tests {
         formatter.write_str("second").unwrap();
         formatter.pop_label().unwrap();
         formatter.write_str("after").unwrap();
+        drop(formatter);
         insta::assert_snapshot!(String::from_utf8(output).unwrap(), @"beforefirstsecondafter");
     }
 
@@ -794,6 +808,7 @@ mod tests {
             .write_str("\x1b[1mnot actually bold\x1b[0m")
             .unwrap();
         formatter.pop_label().unwrap();
+        drop(formatter);
         insta::assert_snapshot!(String::from_utf8(output).unwrap(), @"␛[1mnot actually bold␛[0m");
     }
 
@@ -820,6 +835,7 @@ mod tests {
         formatter.write_str(" after inner ").unwrap();
         formatter.pop_label().unwrap();
         formatter.write_str(" after outer ").unwrap();
+        drop(formatter);
         insta::assert_snapshot!(String::from_utf8(output).unwrap(),
         @" before outer  before inner  inside inner  after inner  after outer ");
     }
@@ -841,6 +857,7 @@ mod tests {
         formatter.pop_label().unwrap();
         formatter.write_str(" not colored ").unwrap();
         formatter.pop_label().unwrap();
+        drop(formatter);
         insta::assert_snapshot!(String::from_utf8(output).unwrap(),
         @" not colored  colored  not colored ");
     }
@@ -885,10 +902,11 @@ mod tests {
         formatter.write_str(" default bg, ").unwrap();
         formatter.pop_label().unwrap();
         formatter.write_str(" and back.").unwrap();
+        drop(formatter);
         insta::assert_snapshot!(String::from_utf8(output).unwrap(),
         @r###"
         Blue on yellow,  default fg,  and back.
-        Blue on yellow,  default bg,  and back.
+        Blue on yellow,  default bg,  and back.
         "###);
     }
 
@@ -908,6 +926,7 @@ mod tests {
         formatter.write_str(" hello ").unwrap();
         formatter.pop_label().unwrap();
         formatter.pop_label().unwrap();
+        drop(formatter);
         insta::assert_snapshot!(String::from_utf8(output).unwrap(),
         @" hello ");
     }
@@ -927,6 +946,7 @@ mod tests {
         formatter.write_str(" hello ").unwrap();
         formatter.pop_label().unwrap();
         formatter.pop_label().unwrap();
+        drop(formatter);
         insta::assert_snapshot!(String::from_utf8(output).unwrap(), @" hello ");
     }
 
@@ -954,10 +974,29 @@ mod tests {
         formatter.pop_label().unwrap();
         formatter.write_str(" a2 ").unwrap();
         formatter.pop_label().unwrap();
+        drop(formatter);
         insta::assert_snapshot!(String::from_utf8(output).unwrap(),
         @" a1  b1  c  b2  a2 ");
     }
 
+    #[test]
+    fn test_color_formatter_dropped() {
+        // Test that the style gets reset if the formatter is dropped without popping
+        // all labels.
+        let config = config_from_string(
+            r#"
+        colors.outer = "green"
+        "#,
+        );
+        let mut output: Vec<u8> = vec![];
+        let mut formatter = ColorFormatter::for_config(&mut output, &config).unwrap();
+        formatter.push_label("outer").unwrap();
+        formatter.push_label("inner").unwrap();
+        formatter.write_str(" inside ").unwrap();
+        drop(formatter);
+        insta::assert_snapshot!(String::from_utf8(output).unwrap(), @" inside ");
+    }
+
     #[test]
     fn test_format_recorder() {
         let mut recorder = FormatRecorder::new();
@@ -977,6 +1016,7 @@ mod tests {
         let mut output: Vec<u8> = vec![];
         let mut formatter = ColorFormatter::for_config(&mut output, &config).unwrap();
         recorder.replay(&mut formatter).unwrap();
+        drop(formatter);
         insta::assert_snapshot!(
             String::from_utf8(output).unwrap(),
             @" outer1  inner1  inner2  outer2 ");
@@ -990,6 +1030,7 @@ mod tests {
                 write!(formatter, "<<{}>>", str::from_utf8(data).unwrap())
             })
             .unwrap();
+        drop(formatter);
         insta::assert_snapshot!(
             String::from_utf8(output).unwrap(),
             @"<< outer1 >><< inner1  inner2 >><< outer2 >>");
diff --git a/cli/src/template_builder.rs b/cli/src/template_builder.rs
index 20539c8c3..d14be4244 100644
--- a/cli/src/template_builder.rs
+++ b/cli/src/template_builder.rs
@@ -946,6 +946,7 @@ mod tests {
             let mut output = Vec::new();
             let mut formatter = ColorFormatter::new(&mut output, self.color_rules.clone().into());
             template.format(&(), &mut formatter).unwrap();
+            drop(formatter);
             String::from_utf8(output).unwrap()
         }
     }
diff --git a/cli/src/text_util.rs b/cli/src/text_util.rs
index fcecef4f8..cefcd1b14 100644
--- a/cli/src/text_util.rs
+++ b/cli/src/text_util.rs
@@ -275,6 +275,7 @@ mod tests {
         let mut output = Vec::new();
         let mut formatter = ColorFormatter::for_config(&mut output, &config).unwrap();
         write(&mut formatter).unwrap();
+        drop(formatter);
         String::from_utf8(output).unwrap()
     }