From 524cd92809add35aad20d3410935374c3b0e4acb Mon Sep 17 00:00:00 2001
From: Stephen <stephen@stephendownward.ca>
Date: Fri, 9 Oct 2020 21:03:15 -0300
Subject: [PATCH] Gracefully handle colours in stdin

---
 Cargo.lock  | 44 ++++++++++++++++++++++++++++++++++++++++++++
 Cargo.toml  |  4 +++-
 src/lol.rs  | 41 ++++++++++++++++++++---------------------
 src/main.rs |  3 +++
 4 files changed, 70 insertions(+), 22 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 604902d..e982336 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1,5 +1,14 @@
 # This file is automatically @generated by Cargo.
 # It is not intended for manual editing.
+[[package]]
+name = "aho-corasick"
+version = "0.7.13"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "043164d8ba5c4c3035fec9bbee8647c0261d788f3474306f93bb65901cae0e86"
+dependencies = [
+ "memchr",
+]
+
 [[package]]
 name = "atty"
 version = "0.2.14"
@@ -78,7 +87,9 @@ version = "0.1.3"
 dependencies = [
  "clap",
  "colored",
+ "lazy_static",
  "rand",
+ "regex",
 ]
 
 [[package]]
@@ -138,6 +149,12 @@ version = "0.2.77"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "f2f96b10ec2560088a8e76961b00d47107b3a625fecb76dedb29ee7ccbf98235"
 
+[[package]]
+name = "memchr"
+version = "2.3.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3728d817d99e5ac407411fa471ff9800a778d88a24685968b36824eaf4bee400"
+
 [[package]]
 name = "os_str_bytes"
 version = "2.3.2"
@@ -233,6 +250,24 @@ dependencies = [
  "rand_core",
 ]
 
+[[package]]
+name = "regex"
+version = "1.3.9"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9c3780fcf44b193bc4d09f36d2a3c87b251da4a046c87795a0d35f4f927ad8e6"
+dependencies = [
+ "aho-corasick",
+ "memchr",
+ "regex-syntax",
+ "thread_local",
+]
+
+[[package]]
+name = "regex-syntax"
+version = "0.6.18"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "26412eb97c6b088a6997e05f69403a802a92d520de2f8e63c2b65f9e0f47c4e8"
+
 [[package]]
 name = "strsim"
 version = "0.10.0"
@@ -268,6 +303,15 @@ dependencies = [
  "unicode-width",
 ]
 
+[[package]]
+name = "thread_local"
+version = "1.0.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d40c6d1b69745a6ec6fb1ca717914848da4b44ae29d9b3080cbee91d72a69b14"
+dependencies = [
+ "lazy_static",
+]
+
 [[package]]
 name = "unicode-segmentation"
 version = "1.6.0"
diff --git a/Cargo.toml b/Cargo.toml
index 2ca9ee3..03da67a 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -16,4 +16,6 @@ categories = ["command-line-utilities"]
 [dependencies]
 colored = "2.0"
 rand = "0.7"
-clap = "3.0.0-beta.2"
\ No newline at end of file
+clap = "3.0.0-beta.2"
+regex = "1"
+lazy_static = "1.4"
\ No newline at end of file
diff --git a/src/lol.rs b/src/lol.rs
index 3dab921..bad8a28 100644
--- a/src/lol.rs
+++ b/src/lol.rs
@@ -1,7 +1,6 @@
-use colored::Colorize;
+use colored::Color;
+use regex::{Regex, RegexBuilder};
 use std::f64::consts;
-use std::io::{stdout, Write};
-use std::process::exit;
 
 fn rainbow(freq: f64, i: f64) -> (u8, u8, u8) {
     let red = ((freq * i).sin() * 127.0) as i16 + 128;
@@ -11,23 +10,23 @@ fn rainbow(freq: f64, i: f64) -> (u8, u8, u8) {
 }
 
 pub fn print_rainbow(line: &str, freq: f64, seed: f64, spread: f64, invert: bool) {
-    let t = line
-        .chars()
-        .enumerate()
-        .map(|(i, c)| {
-            let (r, g, b) = rainbow(freq, seed + i as f64 / spread);
-            if invert {
-                if c == '\n' || c == '\r' {
-                    c.to_string().normal()
-                } else {
-                    c.to_string().on_truecolor(r, g, b)
-                }
-            } else {
-                c.to_string().truecolor(r, g, b)
-            }
-        })
-        .fold("".to_string(), |acc, v| format!("{}{}", acc, v));
-    if stdout().write_all(t.as_bytes()).is_err() {
-        exit(0);
+    lazy_static! {
+        static ref ANSI_ESCAPE: Regex =
+            RegexBuilder::new("((?:\x1B(?:[ -/]+.|[]PX^_][^\x07\x1B]*|\\[[0-?]*.|.))*)(.?)")
+                .multi_line(true)
+                .build()
+                .unwrap();
+    };
+
+    for (i, c) in ANSI_ESCAPE.captures_iter(line).enumerate() {
+        let (r, g, b) = rainbow(freq, seed + i as f64 / spread);
+        let color = Color::TrueColor { r, g, b };
+
+        if invert {
+            print!("{}\x1B[{}m{}\x1B[49m", &c[1], color.to_bg_str(), &c[2]);
+        } else {
+            print!("{}\x1B[{}m{}\x1B[39m", &c[1], color.to_fg_str(), &c[2]);
+        }
     }
+    println!();
 }
diff --git a/src/main.rs b/src/main.rs
index a1c4b13..a5995da 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,3 +1,6 @@
+#[macro_use]
+extern crate lazy_static;
+
 use clap::Clap;
 use std::fs::File;
 use std::io::{stdin, BufRead, BufReader};
-- 
GitLab