Expose bullet_stream::strip_ansi

I already had to write it for the Ruby buildpack (https://github.com/heroku/buildpacks-ruby/blob/c8cdfd0be3a61f7b50d36cae12ec3d22f8068afc/buildpacks/ruby/src/layers/shared.rs#L146), it's useful to have here, and it's useful to show the result of an output in tests like #10.
heroku-buildpacks · Oct 10, 2024 · f1c91af · f1c91af
1 parent 860e573
commit f1c91af
Show file tree

Hide file tree

Showing 3 changed files with 69 additions and 67 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -2,6 +2,8 @@
 
 ## Unreleased
 
+- Expose `bullet_stream::strip_ansi` (https://github.com/schneems/bullet_stream/pull/10)
+
 ## v0.2.0 - 2024/06/06
 
 - Added: `Print` struct. `Output` is now deprecated, use `Print` instead (https://github.com/schneems/bullet_stream/pull/5)

diff --git a/src/ansi_escape.rs b/src/ansi_escape.rs
@@ -58,10 +58,62 @@ impl ANSI {
     }
 }
 
+/// Removes ANSI escape sequences generated from this library from a string.
+///
+/// Not guaranteed to remove all ANSI escape sequences, only those generated by this library.
+pub fn strip_ansi(contents: impl AsRef<str>) -> String {
+    let mut result = String::new();
+    let mut in_sequence = false;
+    for char in contents.as_ref().chars() {
+        // If current character is an escape, set the escape flag which will begin ignoring characters
+        // until the end of the sequence is found.
+        if char == '\x1B' {
+            in_sequence = true;
+        } else if in_sequence {
+            // If we're in a sequence discard the character, an 'm' indicates the end of the sequence
+            if char == 'm' {
+                in_sequence = false;
+            }
+        } else {
+            result.push(char);
+        }
+    }
+
+    result
+}
+
 #[cfg(test)]
 mod test {
     use super::*;
 
+    #[test]
+    fn test_strip_ansi() {
+        for color in [
+            ANSI::Dim,
+            ANSI::Red,
+            ANSI::Yellow,
+            ANSI::BoldCyan,
+            ANSI::BoldPurple,
+            ANSI::BoldUnderlineCyan,
+        ] {
+            // Ensure we covered each color, if a new color is added to ANSI, but not
+            // to this test, it won't compile
+            match &color {
+                ANSI::Dim => (),
+                ANSI::Red => (),
+                ANSI::Yellow => (),
+                ANSI::BoldCyan => (),
+                ANSI::BoldUnderlineCyan => (),
+                ANSI::BoldPurple => (),
+            }
+
+            let input = "Hello world";
+            let output = wrap_ansi_escape_each_line(&color, input);
+
+            assert_eq!(input, &strip_ansi(&output));
+        }
+    }
+
     #[test]
     fn empty_line() {
         let actual = wrap_ansi_escape_each_line(&ANSI::Red, "\n");

diff --git a/src/lib.rs b/src/lib.rs
@@ -16,6 +16,8 @@ pub mod style;
 mod util;
 mod write;
 
+pub use ansi_escape::strip_ansi;
+
 /// Use [`Print`] to output structured text as a buildpack/script executes. The output
 /// is intended to be read by the application user.
 ///
@@ -791,6 +793,7 @@ fn writeln_now<D: Write>(destination: &mut D, msg: impl AsRef<str>) {
 mod test {
     use super::*;
     use crate::util::LockedWriter;
+    use ansi_escape::strip_ansi;
     use fun_run::CommandWithName;
     use indoc::formatdoc;
     use libcnb_test::assert_contains;
@@ -811,10 +814,7 @@ mod test {
             - Done (finished in < 0.1s)
         "};
 
-        assert_eq!(
-            expected,
-            strip_ansi_escape_sequences(String::from_utf8_lossy(&io))
-        )
+        assert_eq!(expected, strip_ansi(String::from_utf8_lossy(&io)))
     }
 
     #[test]
@@ -832,10 +832,7 @@ mod test {
             - Done (finished in < 0.1s)
         "};
 
-        assert_eq!(
-            expected,
-            strip_ansi_escape_sequences(String::from_utf8_lossy(&io))
-        )
+        assert_eq!(expected, strip_ansi(String::from_utf8_lossy(&io)))
     }
 
     #[test]
@@ -866,10 +863,7 @@ mod test {
             - Done (finished in < 0.1s)
         "};
 
-        assert_eq!(
-            expected,
-            strip_ansi_escape_sequences(String::from_utf8_lossy(&io))
-        );
+        assert_eq!(expected, strip_ansi(String::from_utf8_lossy(&io)));
     }
 
     #[test]
@@ -889,10 +883,7 @@ mod test {
             - Done (finished in < 0.1s)
         "};
 
-        assert_eq!(
-            expected,
-            strip_ansi_escape_sequences(String::from_utf8_lossy(&io))
-        );
+        assert_eq!(expected, strip_ansi(String::from_utf8_lossy(&io)));
 
         // Test timer dot colorization
         let expected = formatdoc! {"
@@ -929,10 +920,7 @@ mod test {
             - Done (finished in < 0.1s)
         "};
 
-        assert_eq!(
-            expected,
-            strip_ansi_escape_sequences(String::from_utf8_lossy(&io))
-        );
+        assert_eq!(expected, strip_ansi(String::from_utf8_lossy(&io)));
     }
 
     #[test]
@@ -978,10 +966,7 @@ mod test {
             - Done (finished in < 0.1s)
         "};
 
-        assert_eq!(
-            expected,
-            strip_ansi_escape_sequences(String::from_utf8_lossy(&io))
-        );
+        assert_eq!(expected, strip_ansi(String::from_utf8_lossy(&io)));
     }
 
     #[test]
@@ -1001,10 +986,7 @@ mod test {
 
         "};
 
-        assert_eq!(
-            expected,
-            strip_ansi_escape_sequences(std::fs::read_to_string(path).unwrap())
-        );
+        assert_eq!(expected, strip_ansi(std::fs::read_to_string(path).unwrap()));
     }
 
     #[test]
@@ -1051,10 +1033,7 @@ mod test {
             - Done (finished in < 0.1s)
         "};
 
-        assert_eq!(
-            expected,
-            strip_ansi_escape_sequences(String::from_utf8_lossy(&io))
-        );
+        assert_eq!(expected, strip_ansi(String::from_utf8_lossy(&io)));
     }
 
     #[test]
@@ -1076,7 +1055,7 @@ mod test {
 
         let io = stream.done().done().done();
 
-        let actual = strip_ansi_escape_sequences(String::from_utf8_lossy(&io));
+        let actual = strip_ansi(String::from_utf8_lossy(&io));
 
         assert_contains!(actual, "      hello world\n");
     }
@@ -1106,10 +1085,7 @@ mod test {
             - Done (finished in < 0.1s)
         "};
 
-        assert_eq!(
-            expected,
-            strip_ansi_escape_sequences(String::from_utf8_lossy(&io))
-        );
+        assert_eq!(expected, strip_ansi(String::from_utf8_lossy(&io)));
     }
 
     #[test]
@@ -1140,10 +1116,7 @@ mod test {
             - Done (finished in < 0.1s)
         "};
 
-        assert_eq!(
-            expected,
-            strip_ansi_escape_sequences(String::from_utf8_lossy(&io))
-        );
+        assert_eq!(expected, strip_ansi(String::from_utf8_lossy(&io)));
     }
 
     #[test]
@@ -1178,31 +1151,6 @@ mod test {
             - Done (finished in < 0.1s)
         "};
 
-        assert_eq!(
-            expected,
-            strip_ansi_escape_sequences(String::from_utf8_lossy(&io))
-        );
-    }
-
-    fn strip_ansi_escape_sequences(contents: impl AsRef<str>) -> String {
-        let mut result = String::new();
-        let mut in_ansi_escape = false;
-        for char in contents.as_ref().chars() {
-            if in_ansi_escape {
-                if char == 'm' {
-                    in_ansi_escape = false;
-                    continue;
-                }
-            } else {
-                if char == '\x1B' {
-                    in_ansi_escape = true;
-                    continue;
-                }
-
-                result.push(char);
-            }
-        }
-
-        result
+        assert_eq!(expected, strip_ansi(String::from_utf8_lossy(&io)));
     }
 }