diff --git a/main.go b/main.go index be6ad56..27bacf4 100644 --- a/main.go +++ b/main.go @@ -62,6 +62,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" @@ -80,6 +81,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) } @@ -95,6 +97,7 @@ func main() { versionFlag bool insecureFlag bool jsonFlag bool + filenameFlag bool ) flag.BoolVar(&ungronFlag, "ungron", false, "") @@ -111,6 +114,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() @@ -129,7 +134,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) @@ -168,7 +176,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) @@ -177,6 +190,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