experimental: show context of unfiltered lines

cmd/jlog/main.go

@@ -36,12 +36,16 @@ type output struct {
 	NoSummary          bool     `long:"no-summary" description:"Suppress printing the summary at the end." env:"JLOG_NO_SUMMARY"`
 	PriorityFields     []string `long:"priority" short:"p" description:"A list of fields to show first; repeatable." env:"JLOG_PRIORITY_FIELDS" env-delim:","`
 	HighlightFields    []string `long:"highlight" short:"H" description:"A list of fields to visually distinguish; repeatable." env:"JLOG_HIGHLIGHT_FIELDS" env-delim:"," default:"err" default:"error" default:"warn" default:"warning"`
+
+	AfterContext  int `long:"after-context" short:"A" default:"0" description:"Print this many filtered lines after a non-filtered line (like grep)."`
+	BeforeContext int `long:"before-context" short:"B" default:"0" description:"Print this many filtered lines before a non-filtered line (like grep)."`
+	Context       int `long:"context" short:"C" default:"0" description:"Print this many context lines around each match (like grep)."`
 }
 
 type general struct {
 	JQ           string `short:"e" description:"A jq program to run on each record in the processed input; use this to ignore certain lines, add fields, etc.  Hint: 'select(condition)' will remove lines that don't match 'condition'."`
 	NoColor      bool   `short:"M" long:"no-color" description:"Disable the use of color." env:"JLOG_FORCE_MONOCHROME"`
-	NoMonochrome bool   `short:"C" long:"no-monochrome" description:"Force the use of color.  Note: the short flag will change in a future release." ENV:"JLOG_FORCE_COLOR"`
+	NoMonochrome bool   `short:"c" long:"no-monochrome" description:"Force the use of color." ENV:"JLOG_FORCE_COLOR"`
 	Profile      string `long:"profile" description:"If set, collect a CPU profile and write it to this file."`
 
 	Version bool `short:"v" long:"version" description:"Print version information and exit."`
@@ -220,6 +224,16 @@ func main() {
 	outs := &parse.OutputSchema{
 		Formatter:      defaultOutput,
 		PriorityFields: out.PriorityFields,
+		AfterContext:   out.Context,
+		BeforeContext:  out.Context,
+	}
+
+	// Let -A and -B override -C.
+	if a := outs.AfterContext; a > 0 {
+		outs.AfterContext = a
+	}
+	if b := outs.BeforeContext; b > 0 {
+		outs.BeforeContext = b
 	}
 
 	sigCh := make(chan os.Signal, 1)

pkg/parse/context.go

@@ -0,0 +1,45 @@
+package parse
+
+import (
+	"bytes"
+)
+
+type context struct {
+	Before, After int
+
+	lines      []string
+	printAfter int
+	line       int
+	lastPrint  int
+}
+
+func (c *context) Print(buf *bytes.Buffer, msg string, selected bool) {
+	c.line++
+	if selected {
+		c.printAfter = c.After
+		if c.lastPrint != 0 && (c.After != 0 || c.Before != 0) && c.line-len(c.lines)-c.lastPrint > 1 {
+			buf.WriteString("---\n")
+		}
+		for _, l := range c.lines {
+			buf.WriteString(l)
+		}
+		buf.WriteString(msg)
+		c.lastPrint = c.line
+		c.lines = nil
+		return
+	}
+
+	if c.printAfter > 0 {
+		buf.WriteString(msg)
+		c.lastPrint = c.line
+		c.printAfter--
+		return
+	}
+
+	if c.Before > 0 {
+		c.lines = append(c.lines, msg)
+		if len(c.lines) > c.Before {
+			c.lines = c.lines[1:]
+		}
+	}
+}

pkg/parse/context_test.go

@@ -0,0 +1,160 @@
+package parse
+
+import (
+	"bytes"
+	"regexp"
+	"strings"
+	"testing"
+
+	"github.com/google/go-cmp/cmp"
+	"github.com/google/go-cmp/cmp/cmpopts"
+)
+
+func TestContext(t *testing.T) {
+	testData := []struct {
+		name          string
+		before, after int
+		match         *regexp.Regexp
+		input         []string
+		want          []string
+	}{
+		{
+			name:   "select none",
+			before: 2,
+			after:  2,
+			match:  regexp.MustCompile(`^never matches$`),
+			input:  []string{"1", "2", "3", "4", "5", "6", "7", "8", "9", "10"},
+			want:   []string{},
+		},
+		{
+			name:   "select all",
+			before: 2,
+			after:  2,
+			match:  regexp.MustCompile(`^.*$`),
+			input:  []string{"1", "2", "3", "4", "5", "6", "7", "8", "9", "10"},
+			want:   []string{"1", "2", "3", "4", "5", "6", "7", "8", "9", "10"},
+		},
+		{
+			name:   "no context, single match",
+			before: 0,
+			after:  0,
+			match:  regexp.MustCompile(`^5$`),
+			input:  []string{"1", "2", "3", "4", "5", "6", "7", "8", "9", "10"},
+			want:   []string{"5"},
+		},
+		{
+			name:   "no context, two matches",
+			before: 0,
+			after:  0,
+			match:  regexp.MustCompile(`^5|8$`),
+			input:  []string{"1", "2", "3", "4", "5", "6", "7", "8", "9", "10"},
+			want:   []string{"5", "8"},
+		},
+		{
+			name:   "no context, two contiguous matches",
+			before: 0,
+			after:  0,
+			match:  regexp.MustCompile(`^5|6$`),
+			input:  []string{"1", "2", "3", "4", "5", "6", "7", "8", "9", "10"},
+			want:   []string{"5", "6"},
+		},
+		{
+			name:   "basic context, single match",
+			before: 2,
+			after:  2,
+			match:  regexp.MustCompile(`^5$`),
+			input:  []string{"1", "2", "3", "4", "5", "6", "7", "8", "9", "10"},
+			want:   []string{"3", "4", "5", "6", "7"},
+		},
+		{
+			name:   "basic context, single match, at start",
+			before: 2,
+			after:  2,
+			match:  regexp.MustCompile(`^1$`),
+			input:  []string{"1", "2", "3", "4", "5", "6", "7", "8", "9", "10"},
+			want:   []string{"1", "2", "3"},
+		},
+		{
+			name:   "basic context, single match, at second element",
+			before: 2,
+			after:  2,
+			match:  regexp.MustCompile(`^2$`),
+			input:  []string{"1", "2", "3", "4", "5", "6", "7", "8", "9", "10"},
+			want:   []string{"1", "2", "3", "4"},
+		},
+		{
+			name:   "basic context, separated match regions",
+			before: 1,
+			after:  1,
+			match:  regexp.MustCompile(`^5|9$`),
+			input:  []string{"1", "2", "3", "4", "5", "6", "7", "8", "9", "10"},
+			want:   []string{"4", "5", "6", "---", "8", "9", "10"},
+		},
+		{
+			name:   "basic context, contiguous match regions",
+			before: 1,
+			after:  1,
+			match:  regexp.MustCompile(`^5|8$`),
+			input:  []string{"1", "2", "3", "4", "5", "6", "7", "8", "9", "10"},
+			want:   []string{"4", "5", "6", "7", "8", "9"},
+		},
+		{
+			name:   "after context, separated match regions",
+			before: 0,
+			after:  2,
+			match:  regexp.MustCompile(`^2|8$`),
+			input:  []string{"1", "2", "3", "4", "5", "6", "7", "8", "9", "10"},
+			want:   []string{"2", "3", "4", "---", "8", "9", "10"},
+		},
+		{
+			name:   "after context, contiguous match regions",
+			before: 0,
+			after:  2,
+			match:  regexp.MustCompile(`^2|3$`),
+			input:  []string{"1", "2", "3", "4", "5", "6", "7", "8", "9", "10"},
+			want:   []string{"2", "3", "4", "5"},
+		},
+		{
+			name:   "before context, separated match regions",
+			before: 2,
+			after:  0,
+			match:  regexp.MustCompile(`^4|9$`),
+			input:  []string{"1", "2", "3", "4", "5", "6", "7", "8", "9", "10"},
+			want:   []string{"2", "3", "4", "---", "7", "8", "9"},
+		},
+		{
+			name:   "before context, contiguous match regions",
+			before: 2,
+			after:  0,
+			match:  regexp.MustCompile(`^5|8$`),
+			input:  []string{"1", "2", "3", "4", "5", "6", "7", "8", "9", "10"},
+			want:   []string{"3", "4", "5", "6", "7", "8"},
+		},
+	}
+
+	for _, test := range testData {
+		t.Run(test.name, func(t *testing.T) {
+			ctx := &context{
+				Before: test.before,
+				After:  test.after,
+			}
+			out := new(bytes.Buffer)
+			for _, l := range test.input {
+				selected := test.match.MatchString(l)
+				ctx.Print(out, l+"\n", selected)
+			}
+
+			gotOutput := out.String()
+			var got []string
+			if len(gotOutput) > 0 {
+				got = strings.Split(gotOutput, "\n")
+			}
+			if len(got) > 0 && got[len(got)-1] == "" {
+				got = got[:len(got)-1]
+			}
+			if diff := cmp.Diff(got, test.want, cmpopts.EquateEmpty()); diff != "" {
+				t.Errorf("diff:\n%s", diff)
+			}
+		})
+	}
+}

pkg/parse/parse.go

@@ -97,9 +97,11 @@ type OutputSchema struct {
 	PriorityFields []string         // PriorityFields controls which fields are printed first.
 	Formatter      OutputFormatter  // Actually does the formatting.
 	EmitErrorFn    func(msg string) // A function that sees all errors.
-	state          State            // state carries context between lines
+	BeforeContext  int              // Context lines to print before a match.
+	AfterContext   int              // Context lines to print after a match.
 
 	suppressionConfigured, noTime, noLevel, noMessage bool
+	state                                             State // state carries context between lines
 }
 
 // EmitError prints any internal errors, so that log lines are not silently ignored if they are
@@ -218,12 +220,20 @@ func ReadLog(r io.Reader, w io.Writer, ins *InputSchema, outs *OutputSchema, jq
 		}
 	}
 	var sum Summary
+
+	lineBuf := new(bytes.Buffer)
 	buf := new(bytes.Buffer)
+	ctx := &context{
+		After:  outs.AfterContext,
+		Before: outs.BeforeContext,
+	}
+
 	for s.Scan() {
 		sum.Lines++
 
 		err := func() (retErr error) {
 			var addError, writeRawLine, recoverable bool
+			var filtered bool
 
 			// Adjust counters, print debugging information, flush buffers on the way
 			// out, no matter what.
@@ -232,7 +242,8 @@ func ReadLog(r io.Reader, w io.Writer, ins *InputSchema, outs *OutputSchema, jq
 					sum.Errors++
 				}
 				var writeError bool
-				if buf.Len() > 0 {
+				if lineBuf.Len() > 0 {
+					ctx.Print(buf, lineBuf.String(), !filtered)
 					if _, err := buf.WriteTo(w); err != nil {
 						recoverable = false
 						writeError = true
@@ -276,7 +287,8 @@ func ReadLog(r io.Reader, w io.Writer, ins *InputSchema, outs *OutputSchema, jq
 			}()
 
 			// Reset state from the last line.
-			buf.Truncate(0)
+			buf.Reset()
+			lineBuf.Reset()
 			l.raw = s.Bytes()
 			l.msg = ""
 			l.fields = make(map[string]interface{})
@@ -296,7 +308,8 @@ func ReadLog(r io.Reader, w io.Writer, ins *InputSchema, outs *OutputSchema, jq
 			}
 
 			// Filter.
-			filtered, err := runJQ(jq, &l)
+			var err error
+			filtered, err = runJQ(jq, &l)
 			if err != nil {
 				addError = true
 				writeRawLine = true
@@ -317,17 +330,16 @@ func ReadLog(r io.Reader, w io.Writer, ins *InputSchema, outs *OutputSchema, jq
 					writeRawLine = false
 					return fmt.Errorf("parse: %w", parseErr)
 				}
-				return nil
 			}
 
-			// Emit a line to the output buffer.
+			// Emit a line.
 			if !outs.suppressionConfigured {
 				outs.noTime = ins.NoTimeKey
 				outs.noLevel = ins.NoLevelKey
 				outs.noMessage = ins.NoMessageKey
 				outs.suppressionConfigured = true
 			}
-			outs.Emit(&l, buf)
+			outs.Emit(&l, lineBuf)
 
 			// Copying the buffer to the output writer is handled in defer.
 			if parseErr != nil {