Add JSON stream highlighter

Author: sibprogrammer

cmd/root.go

@@ -72,13 +72,18 @@ func NewRootCmd() *cobra.Command {
 				} else if cssQuery != "" {
 					err = utils.CSSQuery(reader, pw, cssQuery, cssAttr, options)
 				} else {
-					var isHtmlFormatter bool
-					isHtmlFormatter, reader = isHTMLFormatterNeeded(cmd.Flags(), reader)
+					var contentType utils.ContentType
+					contentType, reader = detectFormat(cmd.Flags(), reader)
 
-					if isHtmlFormatter {
+					switch contentType {
+					case utils.ContentHtml:
 						err = utils.FormatHtml(reader, pw, indent, colors)
-					} else {
+					case utils.ContentXml:
 						err = utils.FormatXml(reader, pw, indent, colors)
+					case utils.ContentJson:
+						err = utils.FormatJson(reader, pw, indent, colors)
+					default:
+						err = fmt.Errorf("unknown content type: %v", contentType)
 					}
 				}
 
@@ -198,18 +203,27 @@ func getColorMode(flags *pflag.FlagSet) int {
 	return colors
 }
 
-func isHTMLFormatterNeeded(flags *pflag.FlagSet, origReader io.Reader) (bool, io.Reader) {
+func detectFormat(flags *pflag.FlagSet, origReader io.Reader) (utils.ContentType, io.Reader) {
 	isHtmlFormatter, _ := flags.GetBool("html")
 	if isHtmlFormatter {
-		return isHtmlFormatter, origReader
+		return utils.ContentHtml, origReader
 	}
 
 	buf := make([]byte, 10)
 	length, err := origReader.Read(buf)
 	if err != nil {
-		return false, origReader
+		return utils.ContentText, origReader
 	}
 
 	reader := io.MultiReader(bytes.NewReader(buf[:length]), origReader)
-	return utils.IsHTML(string(buf)), reader
+
+	if utils.IsJSON(string(buf)) {
+		return utils.ContentJson, reader
+	}
+
+	if utils.IsHTML(string(buf)) {
+		return utils.ContentHtml, reader
+	}
+
+	return utils.ContentXml, reader
 }

internal/utils/utils.go

@@ -2,6 +2,7 @@ package utils
 
 import (
 	"bytes"
+	"encoding/json"
 	"encoding/xml"
 	"errors"
 	"fmt"
@@ -14,6 +15,7 @@ import (
 	"io"
 	"os"
 	"os/exec"
+	"reflect"
 	"regexp"
 	"strings"
 )
@@ -24,12 +26,33 @@ const (
 	ColorsDisabled
 )
 
+type ContentType int
+
+const (
+	ContentXml ContentType = iota
+	ContentHtml
+	ContentJson
+	ContentText
+)
+
 type QueryOptions struct {
 	WithTags bool
 	Indent   string
 	Colors   int
 }
 
+const (
+	jsonTokenTopValue = iota
+	jsonTokenArrayStart
+	jsonTokenArrayValue
+	jsonTokenArrayComma
+	jsonTokenObjectStart
+	jsonTokenObjectKey
+	jsonTokenObjectColon
+	jsonTokenObjectValue
+	jsonTokenObjectComma
+)
+
 func FormatXml(reader io.Reader, writer io.Writer, indent string, colors int) error {
 	decoder := xml.NewDecoder(reader)
 	decoder.Strict = false
@@ -344,6 +367,90 @@ func FormatHtml(reader io.Reader, writer io.Writer, indent string, colors int) e
 	return nil
 }
 
+func FormatJson(reader io.Reader, writer io.Writer, indent string, colors int) error {
+	decoder := json.NewDecoder(reader)
+	decoder.UseNumber()
+
+	if ColorsDefault != colors {
+		color.NoColor = colors == ColorsDisabled
+	}
+
+	tagColor := color.New(color.FgYellow).SprintFunc()
+	attrColor := color.New(color.FgHiBlue).SprintFunc()
+	valueColor := color.New(color.FgGreen).SprintFunc()
+
+	level := 0
+	suffix := ""
+	prefix := ""
+
+	for {
+		token, err := decoder.Token()
+
+		if err == io.EOF {
+			break
+		}
+
+		if err != nil {
+			return err
+		}
+
+		v := reflect.ValueOf(*decoder)
+		tokenState := v.FieldByName("tokenState").Int()
+
+		switch tokenState {
+		case jsonTokenObjectColon:
+			suffix = ": "
+		case jsonTokenObjectComma:
+			suffix = ",\n" + strings.Repeat(indent, level)
+		case jsonTokenArrayComma:
+			suffix = ",\n" + strings.Repeat(indent, level)
+		}
+
+		switch tokenType := token.(type) {
+		case json.Delim:
+			switch rune(tokenType) {
+			case '{':
+				_, _ = fmt.Fprintf(writer, "%s%s\n", prefix, tagColor("{"))
+				level++
+				suffix = strings.Repeat(indent, level)
+			case '}':
+				level--
+				_, _ = fmt.Fprintf(writer, "\n%s%s", strings.Repeat(indent, level), tagColor("}"))
+				if tokenState == jsonTokenArrayComma {
+					suffix = ",\n" + strings.Repeat(indent, level)
+				}
+			case '[':
+				_, _ = fmt.Fprintf(writer, "%s%s\n", prefix, tagColor("["))
+				level++
+				suffix = strings.Repeat(indent, level)
+			case ']':
+				level--
+				_, _ = fmt.Fprintf(writer, "\n%s%s", strings.Repeat(indent, level), tagColor("]"))
+			}
+		case string:
+			value := valueColor(token)
+			if tokenState == jsonTokenObjectColon {
+				value = attrColor(token)
+			}
+			_, _ = fmt.Fprintf(writer, "%s\"%s\"", prefix, value)
+		case float64:
+			_, _ = fmt.Fprintf(writer, "%s%v", prefix, valueColor(token))
+		case json.Number:
+			_, _ = fmt.Fprintf(writer, "%s%v", prefix, valueColor(token))
+		case bool:
+			_, _ = fmt.Fprintf(writer, "%s%v", prefix, valueColor(token))
+		case nil:
+			_, _ = fmt.Fprintf(writer, "%s%s", prefix, valueColor("null"))
+		}
+
+		prefix = suffix
+	}
+
+	_, _ = fmt.Fprint(writer, "\n")
+
+	return nil
+}
+
 func IsHTML(input string) bool {
 	input = strings.ToLower(input)
 	htmlMarkers := []string{"html", "<!d", "<body"}
@@ -357,6 +464,12 @@ func IsHTML(input string) bool {
 	return false
 }
 
+func IsJSON(input string) bool {
+	input = strings.ToLower(input)
+	matched, _ := regexp.MatchString(`\s*[{\[]`, input)
+	return matched
+}
+
 func PagerPrint(reader io.Reader, writer io.Writer) error {
 	pager := os.Getenv("PAGER")
 

internal/utils/utils_test.go

@@ -64,9 +64,9 @@ func TestFormatHtml(t *testing.T) {
 	for unformattedFile, expectedFile := range files {
 		unformattedHtmlReader := getFileReader(path.Join("..", "..", "test", "data", "html", unformattedFile))
 
-		bytes, readErr := os.ReadFile(path.Join("..", "..", "test", "data", "html", expectedFile))
+		data, readErr := os.ReadFile(path.Join("..", "..", "test", "data", "html", expectedFile))
 		assert.Nil(t, readErr)
-		expectedHtml := string(bytes)
+		expectedHtml := string(data)
 
 		output := new(strings.Builder)
 		formatErr := FormatHtml(unformattedHtmlReader, output, "  ", ColorsDisabled)
@@ -75,6 +75,25 @@ func TestFormatHtml(t *testing.T) {
 	}
 }
 
+func TestFormatJson(t *testing.T) {
+	files := map[string]string{
+		"unformatted.json": "formatted.json",
+	}
+
+	for unformattedFile, expectedFile := range files {
+		unformattedJsonReader := getFileReader(path.Join("..", "..", "test", "data", "json", unformattedFile))
+
+		data, readErr := os.ReadFile(path.Join("..", "..", "test", "data", "json", expectedFile))
+		assert.Nil(t, readErr)
+		expectedJson := string(data)
+
+		output := new(strings.Builder)
+		formatErr := FormatJson(unformattedJsonReader, output, "  ", ColorsDisabled)
+		assert.Nil(t, formatErr)
+		assert.Equal(t, expectedJson, output.String())
+	}
+}
+
 func TestXPathQuery(t *testing.T) {
 	type test struct {
 		input  string
@@ -134,6 +153,14 @@ func TestIsHTML(t *testing.T) {
 	assert.False(t, IsHTML("<root></root>"))
 }
 
+func TestIsJSON(t *testing.T) {
+	assert.True(t, IsJSON(`{"key": "value"}`))
+	assert.True(t, IsJSON(`{"key": "value", "key2": "value2"}`))
+	assert.True(t, IsJSON(`[1, 2, 3]`))
+	assert.True(t, IsJSON(`   {}`))
+	assert.False(t, IsJSON(`<html></html>`))
+}
+
 func TestPagerPrint(t *testing.T) {
 	var output bytes.Buffer
 	fileReader := getFileReader(path.Join("..", "..", "test", "data", "html", "formatted.html"))

test/data/json/formatted.json

@@ -0,0 +1,25 @@
+{
+  "menu": {
+    "id": "file",
+    "value": 17,
+    "price": 100.32,
+    "popup": {
+      "menuitem": [
+        {
+          "value": "New",
+          "onclick": "CreateNewDoc()"
+        },
+        {
+          "value": "Open",
+          "onclick": "OpenDoc()",
+          "new": true
+        },
+        {
+          "value": "Close",
+          "onclick": "CloseDoc()",
+          "link": null
+        }
+      ]
+    }
+  }
+}

test/data/json/unformatted.json

@@ -0,0 +1,12 @@
+{"menu": {
+  "id": "file" ,
+  "value":  17,
+  "price": 100.32,
+  "popup":    {
+    "menuitem": [
+      {"value": "New", "onclick": "CreateNewDoc()"},
+      {"value": "Open", "onclick": "OpenDoc()", "new": true},
+        {"value": "Close", "onclick": "CloseDoc()", "link": null }
+    ]
+  }
+}}