github.com/ndau/noms@v1.0.5/go/util/writers/max_line_writer.go (about)

     1  // Copyright 2016 Attic Labs, Inc. All rights reserved.
     2  // Licensed under the Apache License, version 2.0:
     3  // http://www.apache.org/licenses/LICENSE-2.0
     4  
     5  package writers
     6  
     7  import "io"
     8  
     9  var (
    10  	// MaxLinesErr is an instance of MaxLinesError that gets returned by
    11  	// Write() whenever the number of lines written has exceeded the number
    12  	// in |MaxLineWriter.MaxLines|.
    13  	MaxLinesErr = MaxLinesError{"Maximum number of lines written"}
    14  )
    15  
    16  // MaxLinesError is the type of error returned by Write() whenever the number
    17  // of lines written has exceeded the number in |MaxLineWriter.MaxLines|.
    18  type MaxLinesError struct {
    19  	msg string
    20  }
    21  
    22  func (e MaxLinesError) Error() string { return e.msg }
    23  
    24  // MaxLineWriter provides an io.Writer interface that counts the number of lines
    25  // that have been written. It will stop writing and returns an error if the
    26  // number of lines written exceeds the number specified in MaxLineWriter.NumLines.
    27  type MaxLineWriter struct {
    28  	Dest     io.Writer
    29  	MaxLines uint32
    30  	NumLines uint32
    31  }
    32  
    33  // Write() stops writing and returns an error if an attempt is made to write
    34  // any byte after |MaxLines| newLines have been written. For example, if MaxLines
    35  // is 1, all bytes will be written up to and including the 1st newline. If there
    36  // are any bytes in |data| after the 1st newline, an error will be returned.
    37  //
    38  // Callers can change the value of |w.MaxLines| before any call to Write().
    39  // Setting MaxLines to 0 will allow any number of newLines.
    40  func (w *MaxLineWriter) Write(data []byte) (int, error) {
    41  	if len(data) == 0 {
    42  		return 0, nil
    43  	}
    44  
    45  	checkMax := w.MaxLines > 0
    46  
    47  	if checkMax && w.NumLines >= w.MaxLines {
    48  		return 0, MaxLinesErr
    49  	}
    50  
    51  	var err error
    52  	byteCnt := 0
    53  
    54  	for i, b := range data {
    55  		if b == byte('\n') {
    56  			w.NumLines++
    57  			if checkMax && w.NumLines > w.MaxLines {
    58  				err = MaxLinesErr
    59  				break
    60  			}
    61  		} else if checkMax && w.NumLines >= w.MaxLines {
    62  			err = MaxLinesErr
    63  			break
    64  		}
    65  		byteCnt = i
    66  	}
    67  
    68  	cnt, err1 := w.Dest.Write(data[:byteCnt+1])
    69  	if err1 != nil {
    70  		return cnt, err1
    71  	}
    72  	return cnt, err
    73  }