From e9c7548b825a4fb56b53111a76cb90e0f6aef4cc Mon Sep 17 00:00:00 2001
From: Stephen Dolan <mu@netsoc.tcd.ie>
Date: Mon, 3 Dec 2012 01:21:07 +0000
Subject: [PATCH] Oh alright then, if you insist.

Colo(u)red output for jq. Enabled by default if isatty(stdout).
Closes #11.
---
 docs/content/3.manual/manual.yml |  6 ++++
 jv.h                             |  2 +-
 jv_print.c                       | 49 +++++++++++++++++++++++++++++---
 main.c                           | 13 +++++++--
 4 files changed, 63 insertions(+), 7 deletions(-)

diff --git a/docs/content/3.manual/manual.yml b/docs/content/3.manual/manual.yml
index 4d228691..0e16a495 100644
--- a/docs/content/3.manual/manual.yml
+++ b/docs/content/3.manual/manual.yml
@@ -65,6 +65,12 @@ sections:
         will result in more compact output by instead putting each
         JSON object on a single line.
 
+      * `--colour-output` / `-C` and `--monochrome-output` / `-M`
+      
+        By default, jq outputs colored JSON if writing to a
+        terminal. You can force it to produce color even if writing to
+        a pipe or a file using `-C`, and disable color with `-M`.
+
       * `--ascii-output` / `-a`
 
         jq usually outputs non-ASCII Unicode codepoints as UTF-8, even
diff --git a/jv.h b/jv.h
index 02817681..12e7c790 100644
--- a/jv.h
+++ b/jv.h
@@ -105,7 +105,7 @@ jv jv_object_iter_value(jv, int);
 
 int jv_get_refcnt(jv);
 
-enum { JV_PRINT_PRETTY = 1, JV_PRINT_ASCII = 2 };
+enum { JV_PRINT_PRETTY = 1, JV_PRINT_ASCII = 2, JV_PRINT_COLOUR = 4 };
 void jv_dump(jv, int flags);
 jv jv_dump_string(jv, int flags);
 
diff --git a/jv_print.c b/jv_print.c
index 2530b23b..acade6cd 100644
--- a/jv_print.c
+++ b/jv_print.c
@@ -6,6 +6,20 @@
 #include "jv_dtoa.h"
 #include "jv_unicode.h"
 
+#define ESC "\033"
+#define COL(c) (ESC "[" c "m")
+#define COLRESET (ESC "[0m")
+
+// Colour table. See http://en.wikipedia.org/wiki/ANSI_escape_code#Colors
+// for how to choose these.
+static jv_kind colour_kinds[] = 
+  {JV_KIND_NULL,   JV_KIND_FALSE, JV_KIND_TRUE, JV_KIND_NUMBER,
+   JV_KIND_STRING, JV_KIND_ARRAY, JV_KIND_OBJECT};
+static const char* colours[] =
+  {COL("30;1"),    COL("0"),      COL("0"),     COL("0"),
+   COL("32"),      COL("37"),     COL("37")};
+#define FIELD_COLOUR COL("34;1")
+
 static void put_buf(const char* s, int len, FILE* fout, jv* strout) {
   if (strout) {
     *strout = jv_string_append_buf(*strout, s, len);
@@ -35,6 +49,7 @@ static void jvp_dump_string(jv str, int ascii_only, FILE* F, jv* S) {
   const char* cstart;
   int c = 0;
   char buf[32];
+  put_char('"', F, S);
   while ((i = jvp_utf8_next((cstart = i), end, &c))) {
     assert(c != -1);
     int unicode_escape = 0;
@@ -91,12 +106,23 @@ static void jvp_dump_string(jv str, int ascii_only, FILE* F, jv* S) {
     }
   }
   assert(c != -1);
+  put_char('"', F, S);
 }
 
 enum { INDENT = 2 };
 
 static void jv_dump_term(struct dtoa_context* C, jv x, int flags, int indent, FILE* F, jv* S) {
   char buf[JVP_DTOA_FMT_MAX_LEN];
+  const char* colour = 0;
+  if (flags & JV_PRINT_COLOUR) {
+    for (unsigned i=0; i<sizeof(colour_kinds)/sizeof(colour_kinds[0]); i++) {
+      if (jv_get_kind(x) == colour_kinds[i]) {
+        colour = colours[i];
+        put_str(colour, F, S);
+        break;
+      }
+    }
+  }
   switch (jv_get_kind(x)) {
   case JV_KIND_INVALID:
     assert(0 && "Invalid value");
@@ -124,9 +150,7 @@ static void jv_dump_term(struct dtoa_context* C, jv x, int flags, int indent, FI
     break;
   }
   case JV_KIND_STRING:
-    put_char('"', F, S);
     jvp_dump_string(x, flags & JV_PRINT_ASCII, F, S);
-    put_char('"', F, S);
     break;
   case JV_KIND_ARRAY: {
     if (jv_array_length(jv_copy(x)) == 0) {
@@ -148,11 +172,13 @@ static void jv_dump_term(struct dtoa_context* C, jv x, int flags, int indent, FI
         }
       }
       jv_dump_term(C, jv_array_get(jv_copy(x), i), flags, indent + INDENT, F, S);
+      if (colour) put_str(colour, F, S);
     }
     if (flags & JV_PRINT_PRETTY) {
       put_char('\n', F, S);
       put_space(indent, F, S);
     }
+    if (colour) put_str(colour, F, S);
     put_char(']', F, S);
     break;
   }
@@ -176,19 +202,34 @@ static void jv_dump_term(struct dtoa_context* C, jv x, int flags, int indent, FI
           put_str(",", F, S);
         }
       }
+      if (colour) put_str(COLRESET, F, S);
+
       first = 0;
-      jv_dump_term(C, jv_object_iter_key(x, i), flags, indent + INDENT, F, S);
-      put_str(":", F, S);
+      jv key = jv_object_iter_key(x, i);
+      if (colour) put_str(FIELD_COLOUR, F, S);
+      jvp_dump_string(key, flags & JV_PRINT_ASCII, F, S);
+      jv_free(key);
+      if (colour) put_str(COLRESET, F, S);
+
+      if (colour) put_str(colour, F, S);
+      put_str((flags & JV_PRINT_PRETTY) ? ": " : ":", F, S);
+      if (colour) put_str(COLRESET, F, S);
+      
       jv_dump_term(C, jv_object_iter_value(x, i), flags, indent + INDENT, F, S);
+      if (colour) put_str(colour, F, S);
     }
     if (flags & JV_PRINT_PRETTY) {
       put_char('\n', F, S);
       put_space(indent, F, S);
     }
+    if (colour) put_str(colour, F, S);
     put_char('}', F, S);
   }
   }
   jv_free(x);
+  if (colour) {
+    put_str(COLRESET, F, S);
+  }
 }
 
 void jv_dump(jv x, int flags) {
diff --git a/main.c b/main.c
index 74ca0e6e..64cb1fd4 100644
--- a/main.c
+++ b/main.c
@@ -2,6 +2,7 @@
 #include <errno.h>
 #include <string.h>
 #include <ctype.h>
+#include <unistd.h>
 #include "compile.h"
 #include "builtin.h"
 #include "jv.h"
@@ -50,8 +51,10 @@ enum {
   RAW_OUTPUT = 8,
   COMPACT_OUTPUT = 16,
   ASCII_OUTPUT = 32,
+  COLOUR_OUTPUT = 64,
+  NO_COLOUR_OUTPUT = 128,
 
-  FROM_FILE = 64,
+  FROM_FILE = 256,
 };
 static int options = 0;
 static struct bytecode* bc;
@@ -63,9 +66,11 @@ static void process(jv value) {
     if ((options & RAW_OUTPUT) && jv_get_kind(result) == JV_KIND_STRING) {
       fwrite(jv_string_value(result), 1, jv_string_length(jv_copy(result)), stdout);
     } else {
-      int dumpopts = 0;
+      int dumpopts = isatty(fileno(stdout)) ? JV_PRINT_COLOUR : 0;
       if (!(options & COMPACT_OUTPUT)) dumpopts |= JV_PRINT_PRETTY;
       if (options & ASCII_OUTPUT) dumpopts |= JV_PRINT_ASCII;
+      if (options & COLOUR_OUTPUT) dumpopts |= JV_PRINT_COLOUR;
+      if (options & NO_COLOUR_OUTPUT) dumpopts &= ~JV_PRINT_COLOUR;
       jv_dump(result, dumpopts);
     }
     printf("\n");
@@ -111,6 +116,10 @@ int main(int argc, char* argv[]) {
       options |= RAW_OUTPUT;
     } else if (isoption(argv[i], 'c', "compact-output")) {
       options |= COMPACT_OUTPUT;
+    } else if (isoption(argv[i], 'C', "color-output")) {
+      options |= COLOUR_OUTPUT;
+    } else if (isoption(argv[i], 'M', "monochrome-output")) {
+      options |= NO_COLOUR_OUTPUT;
     } else if (isoption(argv[i], 'a', "ascii-output")) {
       options |= ASCII_OUTPUT;
     } else if (isoption(argv[i], 'R', "raw-input")) {