Skip to content

Commit

Permalink
Update First(n int) to read no more data than necessary (#199)
Browse files Browse the repository at this point in the history
First(N) now closes its output after it has read N lines, terminating the pipe 
without consuming unnecessary further input (apart from the 4KiB scanner 
buffer).

Co-authored-by: John Arundel <john@bitfieldconsulting.com>
  • Loading branch information
bartdeboer and bitfield committed May 6, 2024
1 parent cfa12bb commit 507316b
Show file tree
Hide file tree
Showing 2 changed files with 28 additions and 7 deletions.
17 changes: 10 additions & 7 deletions script.go
Original file line number Diff line number Diff line change
Expand Up @@ -514,21 +514,24 @@ func (p *Pipe) FilterScan(filter func(string, io.Writer)) *Pipe {

// First produces only the first n lines of the pipe's contents, or all the
// lines if there are less than n. If n is zero or negative, there is no output
// at all.
// at all. When n lines have been produced, First stops reading its input and
// sends EOF to its output.
func (p *Pipe) First(n int) *Pipe {
if p.Error() != nil {
return p
}
if n <= 0 {
return NewPipe()
}
i := 0
return p.FilterScan(func(line string, w io.Writer) {
if i >= n {
return
return p.Filter(func(r io.Reader, w io.Writer) error {
scanner := newScanner(r)
for i := 0; i < n && scanner.Scan(); i++ {
_, err := fmt.Fprintln(w, scanner.Text())
if err != nil {
return err
}
}
fmt.Fprintln(w, line)
i++
return scanner.Err()
})
}

Expand Down
18 changes: 18 additions & 0 deletions script_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -573,6 +573,24 @@ func TestFirstHasNoEffectGivenLessThanNInputLines(t *testing.T) {
}
}

func TestFirstDoesNotConsumeUnnecessaryData(t *testing.T) {
t.Parallel()
// First uses a 4096-byte buffer, so will always read at least
// that much, but no more (once N lines have been read).
r := strings.NewReader(strings.Repeat("line\n", 1000))
got, err := script.NewPipe().WithReader(r).First(1).String()
if err != nil {
t.Fatal(err)
}
want := "line\n"
if want != got {
t.Errorf("want output %q, got %q", want, got)
}
if r.Len() == 0 {
t.Errorf("no data left in reader")
}
}

func TestFreqHandlesLongLines(t *testing.T) {
t.Parallel()
got, err := script.Echo(longLine).Freq().Slice()
Expand Down

0 comments on commit 507316b

Please sign in to comment.