oss.indeed.com/go/go-opine@v1.3.0/internal/printing/line_prefix_writer.go (about)

     1  package printing
     2  
     3  import (
     4  	"bytes"
     5  	"io"
     6  )
     7  
     8  // NewLinePrefixWriter creates and returns a new LinePrefixWriter.
     9  func NewLinePrefixWriter(to io.Writer, prefix string) *LinePrefixWriter {
    10  	return &LinePrefixWriter{
    11  		to:     to,
    12  		prefix: []byte(prefix),
    13  	}
    14  }
    15  
    16  // NewLinePrefixWriter wraps an io.Writer and adds a prefix to every line.
    17  type LinePrefixWriter struct {
    18  	to                 io.Writer
    19  	prefix             []byte
    20  	suppressNextPrefix bool
    21  }
    22  
    23  var _ io.Writer = (*LinePrefixWriter)(nil)
    24  
    25  // Write to the underlying io.Writer, prefixing each new line with the
    26  // configured prefix. Returns the number of *input* bytes written (i.e.
    27  // not including the prefix bytes written) and whether an error occurred.
    28  // If no error occurred the number of bytes written will always be equal
    29  // to the length of the input.
    30  func (w *LinePrefixWriter) Write(p []byte) (int, error) {
    31  	if len(p) == 0 {
    32  		return 0, nil
    33  	}
    34  
    35  	n := 0
    36  	rest := p
    37  	for len(rest) > 0 {
    38  		var line []byte
    39  
    40  		if i := bytes.IndexByte(rest, '\n'); i >= 0 {
    41  			line = rest[:i+1]
    42  			rest = rest[i+1:]
    43  		} else {
    44  			line = rest
    45  			rest = nil
    46  		}
    47  
    48  		var prefix []byte
    49  		if !w.suppressNextPrefix {
    50  			prefix = w.prefix
    51  		}
    52  
    53  		if cnt, err := w.to.Write(append(prefix, line...)); err != nil {
    54  			if cnt < len(prefix) {
    55  				cnt = 0
    56  			} else {
    57  				cnt -= len(prefix)
    58  			}
    59  			return n + cnt, err
    60  		}
    61  		n += len(line)
    62  		w.suppressNextPrefix = false
    63  	}
    64  
    65  	w.suppressNextPrefix = !bytes.HasSuffix(p, []byte{'\n'})
    66  
    67  	return n, nil
    68  }