Add color option

src/main.rs

@@ -8,7 +8,7 @@
   Run `jsonxf -h` for usage options.
 */
 
-use std::{fs::File, io::ErrorKind};
+use std::{fs::File, io::{ErrorKind, Write}};
 
 extern crate jsonxf;
 
@@ -58,6 +58,11 @@ fn do_main() -> Result<(), String> {
         "minimize",
         "minimize JSON instead of pretty-printing it",
     );
+    opts.optflag(
+        "c",
+        "color",
+        "colorize output (like jq)",
+    );
     opts.optflag("h", "help", "print this message and exit");
 
     let matches = match opts.parse(&args[1..]) {
@@ -106,7 +111,7 @@ fn do_main() -> Result<(), String> {
         }
     };
 
-    let mut output: Box<dyn std::io::Write> = match matches.opt_str("o") {
+    let raw_output: Box<dyn std::io::Write> = match matches.opt_str("o") {
         None => Box::new(std::io::stdout()),
         Some(filename) => {
             if filename == *"-" {
@@ -144,6 +149,12 @@ fn do_main() -> Result<(), String> {
         Some(string) => string,
     };
 
+    let mut output: Box<dyn std::io::Write> = if matches.opt_present("c") {
+        Box::new(ColorWriter::new(raw_output))
+    } else {
+        raw_output
+    };
+
     let result = if matches.opt_present("m") {
         let mut xf = jsonxf::Formatter::minimizer();
         xf.format_stream(&mut input, &mut output)
@@ -167,6 +178,128 @@ fn do_main() -> Result<(), String> {
     }
 }
 
+const RESET: &[u8] = b"\x1b[0m";
+const KEY: &[u8] = b"\x1b[1;34m";   // bold blue
+const STRING: &[u8] = b"\x1b[32m";  // green
+const NUMBER: &[u8] = b"\x1b[33m";  // yellow
+const LITERAL: &[u8] = b"\x1b[1m";  // bold (true/false/null)
+
+enum ColorState {
+    Idle,
+    InString { is_key: bool },
+    InStringEscape { is_key: bool },
+    InLiteral,
+    InNumber,
+}
+
+struct ColorWriter<W: Write> {
+    inner: W,
+    state: ColorState,
+    stack: Vec<bool>, // true = object (next string is a key)
+    expect_key: bool,
+}
+
+impl<W: Write> ColorWriter<W> {
+    fn new(inner: W) -> Self {
+        ColorWriter { inner, state: ColorState::Idle, stack: Vec::new(), expect_key: false }
+    }
+
+    fn process_byte(&mut self, b: u8) -> std::io::Result<()> {
+        match self.state {
+            ColorState::Idle => match b {
+                b'{' => {
+                    self.stack.push(true);
+                    self.expect_key = true;
+                    self.inner.write_all(&[b])
+                }
+                b'}' => {
+                    self.stack.pop();
+                    self.expect_key = *self.stack.last().unwrap_or(&false);
+                    self.inner.write_all(&[b])
+                }
+                b'[' => {
+                    self.stack.push(false);
+                    self.expect_key = false;
+                    self.inner.write_all(&[b])
+                }
+                b']' => {
+                    self.stack.pop();
+                    self.expect_key = false;
+                    self.inner.write_all(&[b])
+                }
+                b',' => {
+                    self.expect_key = *self.stack.last().unwrap_or(&false);
+                    self.inner.write_all(&[b])
+                }
+                b'"' => {
+                    let is_key = self.expect_key;
+                    self.inner.write_all(if is_key { KEY } else { STRING })?;
+                    self.inner.write_all(&[b])?;
+                    self.state = ColorState::InString { is_key };
+                    Ok(())
+                }
+                b't' | b'f' | b'n' => {
+                    self.inner.write_all(LITERAL)?;
+                    self.inner.write_all(&[b])?;
+                    self.state = ColorState::InLiteral;
+                    Ok(())
+                }
+                b'0'..=b'9' | b'-' => {
+                    self.inner.write_all(NUMBER)?;
+                    self.inner.write_all(&[b])?;
+                    self.state = ColorState::InNumber;
+                    Ok(())
+                }
+                _ => self.inner.write_all(&[b]),
+            },
+            ColorState::InString { is_key } => match b {
+                b'\\' => {
+                    self.inner.write_all(&[b])?;
+                    self.state = ColorState::InStringEscape { is_key };
+                    Ok(())
+                }
+                b'"' => {
+                    self.inner.write_all(&[b])?;
+                    self.inner.write_all(RESET)?;
+                    if is_key {
+                        self.expect_key = false;
+                    }
+                    self.state = ColorState::Idle;
+                    Ok(())
+                }
+                _ => self.inner.write_all(&[b]),
+            },
+            ColorState::InStringEscape { is_key } => {
+                self.inner.write_all(&[b])?;
+                self.state = ColorState::InString { is_key };
+                Ok(())
+            }
+            ColorState::InLiteral | ColorState::InNumber => {
+                if matches!(b, b' ' | b'\n' | b'\r' | b'\t' | b',' | b'}' | b']' | b':') {
+                    self.inner.write_all(RESET)?;
+                    self.state = ColorState::Idle;
+                    self.process_byte(b)
+                } else {
+                    self.inner.write_all(&[b])
+                }
+            }
+        }
+    }
+}
+
+impl<W: Write> Write for ColorWriter<W> {
+    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
+        for &b in buf {
+            self.process_byte(b)?;
+        }
+        Ok(buf.len())
+    }
+
+    fn flush(&mut self) -> std::io::Result<()> {
+        self.inner.flush()
+    }
+}
+
 fn print_help(program_name: &str, opts: &Options) {
     let desc = "Jsonxf is a JSON transformer.  It provides fast pretty-printing and
 minimizing of JSON-encoded UTF-8 data.";