From 3ac5c0b7eff2453994a5a0c92bc614f0b3a12318 Mon Sep 17 00:00:00 2001
From: Ross Patterson <rosspatt@amazon.com>
Date: Fri, 26 Jun 2020 17:27:56 +0000
Subject: [PATCH] Add a -H/--with-filename option ala grep, to include the
 filename in gron'ed output, for when processing multiple files (e.g., within
 find).

---
 main.go      | 33 +++++++++++++++++++++++++++++++--
 main_test.go | 38 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 69 insertions(+), 2 deletions(-)

diff --git a/main.go b/main.go
index 840e36a..1081072 100644
--- a/main.go
+++ b/main.go
@@ -61,6 +61,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 += "  -H, --with-filename Print the filename on each line\n"
 		h += "      --no-sort    Don't sort output (faster)\n"
 		h += "      --version    Print version information\n\n"
 
@@ -79,6 +80,7 @@ func init() {
 		h += "  gron http://jsonplaceholder.typicode.com/users/1 \n"
 		h += "  curl -s http://jsonplaceholder.typicode.com/users/1 | gron\n"
 		h += "  gron http://jsonplaceholder.typicode.com/users/1 | grep company | gron --ungron\n"
+		h += "  find . -type f -exec gron -H ; | grep banana\n"
 
 		fmt.Fprintf(os.Stderr, h)
 	}
@@ -94,6 +96,7 @@ func main() {
 		versionFlag    bool
 		insecureFlag   bool
 		jsonFlag       bool
+		filenameFlag   bool
 	)
 
 	flag.BoolVar(&ungronFlag, "ungron", false, "")
@@ -110,6 +113,8 @@ func main() {
 	flag.BoolVar(&insecureFlag, "insecure", false, "")
 	flag.BoolVar(&jsonFlag, "j", false, "")
 	flag.BoolVar(&jsonFlag, "json", false, "")
+	flag.BoolVar(&filenameFlag, "H", false, "")
+	flag.BoolVar(&filenameFlag, "with-filename", false, "")
 
 	flag.Parse()
 
@@ -123,7 +128,10 @@ func main() {
 	// file, HTTP URL or stdin
 	var rawInput io.Reader
 	filename := flag.Arg(0)
-	if filename == "" || filename == "-" {
+	if filename == "" {
+		filename = "-"
+	}
+	if filename == "-" {
 		rawInput = os.Stdin
 	} else if validURL(filename) {
 		r, err := getURL(filename, insecureFlag)
@@ -162,7 +170,12 @@ func main() {
 	} else if streamFlag {
 		a = gronStream
 	}
-	exitCode, err := a(rawInput, colorable.NewColorableStdout(), opts)
+	var w io.Writer
+	w = colorable.NewColorableStdout()
+	if filenameFlag {
+		w = NewPrefixedWriter(w, filename + ": ")
+	}
+	exitCode, err := a(rawInput, w, opts)
 
 	if exitCode != exitOK {
 		fatal(exitCode, err)
@@ -171,6 +184,22 @@ func main() {
 	os.Exit(exitOK)
 }
 
+// A prefixedWriter pre-pends an initially-specified string to every Write()
+// call.  It wraps another writer, which is used to deliver the combined string
+// to wherever it belongs.
+type prefixedWriter struct {
+	wrappedWriter io.Writer
+	prefixBytes   []byte
+}
+
+func NewPrefixedWriter(w io.Writer, prefix string) io.Writer {
+	return &prefixedWriter{wrappedWriter: w, prefixBytes: []byte(prefix)}
+}
+
+func (w prefixedWriter) Write(p []byte) (int, error) {
+	return w.wrappedWriter.Write(append(w.prefixBytes, p...))
+}
+
 // an actionFn represents a main action of the program, it accepts
 // an input, output and a bitfield of options; returning an exit
 // code and any error that occurred
diff --git a/main_test.go b/main_test.go
index 94b57cd..182d21b 100644
--- a/main_test.go
+++ b/main_test.go
@@ -1,8 +1,10 @@
 package main
 
 import (
+	"bufio"
 	"bytes"
 	"encoding/json"
+	"fmt"
 	"io/ioutil"
 	"os"
 	"reflect"
@@ -50,6 +52,42 @@ func TestGron(t *testing.T) {
 
 }
 
+func TestPrefixedFilenameWriter(t *testing.T) {
+	prefixes := []string{
+		"testdata/one.json: ",
+		"testdata/two.json: ",
+		"testdata/three.json: ",
+		"testdata/github.json: ",
+	}
+	lines := []string{
+		"Line 1",
+		"Another line 2",
+		"a 3rd Line",
+	}
+
+	for _, p := range prefixes {
+		wrappedWriter := &bytes.Buffer{}
+		writer := NewPrefixedWriter(wrappedWriter, p)
+		for _, l := range lines {
+			fmt.Fprintln(writer, l)
+		}
+		s := bufio.NewScanner(wrappedWriter)
+		var have []string
+		for s.Scan() {
+			have = append(have, s.Text())
+		}
+		if len(have) != len(lines) {
+			t.Errorf("want %d lines, have %d lines", len(lines), len(have))
+		}
+		for i, l := range lines {
+			want := p + l
+			if have[i] != want {
+				t.Errorf("line %d\nwant: %s\nhave: %s", i, want, have[i])
+			}
+		}
+	}
+}
+
 func TestGronStream(t *testing.T) {
 	cases := []struct {
 		inFile  string