Make it possible to process yaml instead of json

go.mod

@@ -1,9 +1,14 @@
 module github.com/tomnomnom/gron
 
+go 1.14
+
 require (
 	github.com/fatih/color v1.7.0
+	github.com/go-yaml/yaml v2.1.0+incompatible
 	github.com/mattn/go-colorable v0.0.9
 	github.com/mattn/go-isatty v0.0.4 // indirect
 	github.com/nwidger/jsoncolor v0.0.0-20170215171346-75a6de4340e5
 	github.com/pkg/errors v0.8.0
+	golang.org/x/sys v0.0.0-20200831180312-196b9ba8737a // indirect
+	gopkg.in/yaml.v2 v2.3.0 // indirect
 )

go.sum

@@ -1,5 +1,7 @@
 github.com/fatih/color v1.7.0 h1:DkWD4oS2D8LGGgTQ6IvwJJXSL5Vp2ffcQg58nFV38Ys=
 github.com/fatih/color v1.7.0/go.mod h1:Zm6kSWBoL9eyXnKyktHP6abPY2pDugNf5KwzbycvMj4=
+github.com/go-yaml/yaml v2.1.0+incompatible h1:RYi2hDdss1u4YE7GwixGzWwVo47T8UQwnTLB6vQiq+o=
+github.com/go-yaml/yaml v2.1.0+incompatible/go.mod h1:w2MrLa16VYP0jy6N7M5kHaCkaLENm+P+Tv+MfurjSw0=
 github.com/mattn/go-colorable v0.0.9 h1:UVL0vNpWh04HeJXV0KLcaT7r06gOH2l4OW6ddYRUIY4=
 github.com/mattn/go-colorable v0.0.9/go.mod h1:9vuHe8Xs5qXnSaW/c/ABM9alt+Vo+STaOChaDxuIBZU=
 github.com/mattn/go-isatty v0.0.4 h1:bnP0vzxcAdeI1zdubAl5PjU6zsERjGZb7raWodagDYs=
@@ -8,3 +10,9 @@ github.com/nwidger/jsoncolor v0.0.0-20170215171346-75a6de4340e5 h1:d+C3xJdxZT7wN
 github.com/nwidger/jsoncolor v0.0.0-20170215171346-75a6de4340e5/go.mod h1:GYFm0zZgTNeoK1QxuIofRDasy2ibmaJZhZLzwsMXUF4=
 github.com/pkg/errors v0.8.0 h1:WdK/asTD0HN+q6hsWO3/vpuAkAr+tw6aNJNDFFf0+qw=
 github.com/pkg/errors v0.8.0/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
+golang.org/x/sys v0.0.0-20200831180312-196b9ba8737a h1:i47hUS795cOydZI4AwJQCKXOr4BvxzvikwDoDtHhP2Y=
+golang.org/x/sys v0.0.0-20200831180312-196b9ba8737a/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/yaml.v2 v2.3.0 h1:clyUAQHOM3G0M3f5vQj7LuJrETvjVot3Z5el9nffUtU=
+gopkg.in/yaml.v2 v2.3.0/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=

main.go

@@ -11,6 +11,7 @@ import (
 	"sort"
 
 	"github.com/fatih/color"
+	"github.com/go-yaml/yaml"
 	"github.com/mattn/go-colorable"
 	"github.com/nwidger/jsoncolor"
 	"github.com/pkg/errors"
@@ -32,6 +33,7 @@ const (
 	optMonochrome = 1 << iota
 	optNoSort
 	optJSON
+	optYAML
 )
 
 // Output colors
@@ -49,7 +51,7 @@ var gronVersion = "dev"
 
 func init() {
 	flag.Usage = func() {
-		h := "Transform JSON (from a file, URL, or stdin) into discrete assignments to make it greppable\n\n"
+		h := "Transform JSON or YAML (from a file, URL, or stdin) into discrete assignments to make it greppable\n\n"
 
 		h += "Usage:\n"
 		h += "  gron [OPTIONS] [FILE|URL|-]\n\n"
@@ -61,6 +63,7 @@ func init() {
 		h += "  -s, --stream     Treat each line of input as a separate JSON object\n"
 		h += "  -k, --insecure   Disable certificate validation\n"
 		h += "  -j, --json       Represent gron data as JSON stream\n"
+		h += "  -y, --yaml       Treat input as YAML instead of JSON\n"
 		h += "      --no-sort    Don't sort output (faster)\n"
 		h += "      --version    Print version information\n\n"
 
@@ -94,6 +97,7 @@ func main() {
 		versionFlag    bool
 		insecureFlag   bool
 		jsonFlag       bool
+		yamlFlag       bool
 	)
 
 	flag.BoolVar(&ungronFlag, "ungron", false, "")
@@ -110,6 +114,8 @@ func main() {
 	flag.BoolVar(&insecureFlag, "insecure", false, "")
 	flag.BoolVar(&jsonFlag, "j", false, "")
 	flag.BoolVar(&jsonFlag, "json", false, "")
+	flag.BoolVar(&yamlFlag, "y", false, "")
+	flag.BoolVar(&yamlFlag, "yaml", false, "")
 
 	flag.Parse()
 
@@ -154,6 +160,9 @@ func main() {
 	if jsonFlag {
 		opts = opts | optJSON
 	}
+	if yamlFlag {
+		opts = opts | optYAML
+	}
 
 	// Pick the appropriate action: gron, ungron or gronStream
 	var a actionFn = gron
@@ -176,6 +185,20 @@ func main() {
 // code and any error that occurred
 type actionFn func(io.Reader, io.Writer, int) (int, error)
 
+type decoder interface {
+	Decode(interface{}) error
+}
+
+func makeDecoder(r io.Reader, opts int) decoder {
+	if opts&optYAML > 0 {
+		return yaml.NewDecoder(r)
+	} else {
+		d := json.NewDecoder(r)
+		d.UseNumber()
+		return d
+	}
+}
+
 // gron is the default action. Given JSON as the input it returns a list
 // of assignment statements. Possible options are optNoSort and optMonochrome
 func gron(r io.Reader, w io.Writer, opts int) (int, error) {
@@ -188,7 +211,12 @@ func gron(r io.Reader, w io.Writer, opts int) (int, error) {
 		conv = statementToColorString
 	}
 
-	ss, err := statementsFromJSON(r, statement{{"json", typBare}})
+	top := "json"
+	if opts&optYAML > 0 {
+		top = "yaml"
+	}
+
+	ss, err := statementsFromJSON(makeDecoder(r, opts), statement{{top, typBare}})
 	if err != nil {
 		goto out
 	}
@@ -268,10 +296,10 @@ func gronStream(r io.Reader, w io.Writer, opts int) (int, error) {
 	i = 0
 	for sc.Scan() {
 
-		line := bytes.NewBuffer(sc.Bytes())
+		d := makeDecoder(bytes.NewBuffer(sc.Bytes()), opts)
 
 		var ss statements
-		ss, err = statementsFromJSON(line, makePrefix(i))
+		ss, err = statementsFromJSON(d, makePrefix(i))
 		i++
 		if err != nil {
 			goto out

statements.go

@@ -3,7 +3,6 @@ package main
 import (
 	"encoding/json"
 	"fmt"
-	"io"
 	"reflect"
 	"strconv"
 	"strings"
@@ -385,11 +384,9 @@ func (ss statements) Contains(search statement) bool {
 
 // statementsFromJSON takes an io.Reader containing JSON
 // and returns statements or an error on failure
-func statementsFromJSON(r io.Reader, prefix statement) (statements, error) {
+func statementsFromJSON(r decoder, prefix statement) (statements, error) {
 	var top interface{}
-	d := json.NewDecoder(r)
-	d.UseNumber()
-	err := d.Decode(&top)
+	err := r.Decode(&top)
 	if err != nil {
 		return nil, err
 	}
@@ -408,6 +405,16 @@ func (ss *statements) fill(prefix statement, v interface{}) {
 	// Recurse into objects and arrays
 	switch vv := v.(type) {
 
+	case map[interface{}]interface{}:
+		// It's an object
+		for k, sub := range vv {
+			ks := fmt.Sprintf("%v", k)
+			if validIdentifier(ks) {
+				ss.fill(prefix.withBare(ks), sub)
+			} else {
+				ss.fill(prefix.withQuotedKey(ks), sub)
+			}
+		}
 	case map[string]interface{}:
 		// It's an object
 		for k, sub := range vv {

statements_test.go

@@ -33,7 +33,7 @@ func TestStatementsSimple(t *testing.T) {
 		"": 2
 	}`)
 
-	ss, err := statementsFromJSON(bytes.NewReader(j), statement{{"json", typBare}})
+	ss, err := statementsFromJSON(makeDecoder(bytes.NewReader(j), 0), statement{{"json", typBare}})
 
 	if err != nil {
 		t.Errorf("Want nil error from makeStatementsFromJSON() but got %s", err)
@@ -65,6 +65,56 @@ func TestStatementsSimple(t *testing.T) {
 
 }
 
+func TestStatementsSimpleYaml(t *testing.T) {
+
+	j := []byte(`'': 2
+a quoted: value
+anarr:
+- 1
+- 1.5
+anob:
+  foo: bar
+anull: 
+bool1: true
+bool2: false
+dotted: A dotted value
+else: 1
+x: |
+  y: "z"
+id: 66912849`)
+
+	ss, err := statementsFromJSON(makeDecoder(bytes.NewReader(j), optYAML), statement{{"yaml", typBare}})
+
+	if err != nil {
+		t.Errorf("Want nil error from makeStatementsFromJSON() but got %s", err)
+	}
+
+	wants := statementsFromStringSlice([]string{
+		`yaml = {};`,
+		`yaml.dotted = "A dotted value";`,
+		`yaml["a quoted"] = "value";`,
+		`yaml.bool1 = true;`,
+		`yaml.bool2 = false;`,
+		`yaml.anull = null;`,
+		`yaml.anarr = [];`,
+		`yaml.anarr[0] = 1;`,
+		`yaml.anarr[1] = 1.5;`,
+		`yaml.anob = {};`,
+		`yaml.anob.foo = "bar";`,
+		`yaml["else"] = 1;`,
+		`yaml.id = 66912849;`,
+		`yaml[""] = 2;`,
+		`yaml.x = "y: \"z\"\n";`,
+	})
+
+	t.Logf("Have: %#v", ss)
+	for _, want := range wants {
+		if !ss.Contains(want) {
+			t.Errorf("Statement group should contain `%s` but doesn't", want)
+		}
+	}
+
+}
 func TestStatementsSorting(t *testing.T) {
 	want := statementsFromStringSlice([]string{
 		`json.a = true;`,

token.go

@@ -116,10 +116,14 @@ func (t token) formatColor() string {
 func valueTokenFromInterface(v interface{}) token {
 	switch vv := v.(type) {
 
+	case map[interface{}]interface{}:
+		return token{"{}", typEmptyObject}
 	case map[string]interface{}:
 		return token{"{}", typEmptyObject}
 	case []interface{}:
 		return token{"[]", typEmptyArray}
+	case int, float64:
+		return token{fmt.Sprintf("%v", vv), typNumber}
 	case json.Number:
 		return token{vv.String(), typNumber}
 	case string: