github.com/attic-labs/noms@v0.0.0-20210827224422-e5fa29d95e8b/go/util/writers/prefix_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  // PrefixWriter makes it easy to prefix lines with a custom prefix. Each time
    10  // it writes a byte after a newline('\n') character it calls PrefixFunc() to get
    11  // the byte slice that should be written. |NeedsPrefix| can be set to true to
    12  // cause a prefix to be written immediately. This is useful for causing a prefix
    13  // to get written on the first line.
    14  type PrefixWriter struct {
    15  	Dest        io.Writer
    16  	PrefixFunc  func(w *PrefixWriter) []byte
    17  	NeedsPrefix bool
    18  	NumLines    uint32
    19  }
    20  
    21  // Write() will add a prefix to the beginning of each line. It obtains the
    22  // prefix by call |PrefixFunc(w *PrefixWriter)| before printing out any character
    23  // following a newLine. Callers can force a prefix to be printed out before the
    24  // first character in |data| by setting NeedsPrefix to true. Conversely, callers
    25  // can suppress prefixes from being printed by setting NeedsPrefix to false.
    26  
    27  func (w *PrefixWriter) Write(data []byte) (int, error) {
    28  	writtenCnt := 0
    29  	for i, b := range data {
    30  		if w.NeedsPrefix {
    31  			w.NeedsPrefix = false
    32  			d1 := w.PrefixFunc(w)
    33  			cnt, err := w.Dest.Write(d1)
    34  			writtenCnt += cnt
    35  			if err != nil {
    36  				return writtenCnt, err
    37  			}
    38  		}
    39  		if b == byte('\n') {
    40  			w.NumLines++
    41  			w.NeedsPrefix = true
    42  		}
    43  		cnt, err := w.Dest.Write(data[i : i+1])
    44  		writtenCnt += cnt
    45  		if err != nil {
    46  			return writtenCnt, err
    47  		}
    48  	}
    49  	return writtenCnt, nil
    50  }