diff --git a/src/main.rs b/src/main.rs index 4038ef2..be670ba 100644 --- a/src/main.rs +++ b/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 = match matches.opt_str("o") { + let raw_output: Box = 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 = 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 { + inner: W, + state: ColorState, + stack: Vec, // true = object (next string is a key) + expect_key: bool, +} + +impl ColorWriter { + 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 Write for ColorWriter { + fn write(&mut self, buf: &[u8]) -> std::io::Result { + 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.";