github.com/grailbio/base@v0.0.11/iofmt/prefixwriter.go (about)

     1  // Copyright 2018 GRAIL, Inc. All rights reserved.
     2  // Use of this source code is governed by the Apache 2.0
     3  // license that can be found in the LICENSE file.
     4  
     5  // Package iofmt contains various IO utilities used for formatting
     6  // output.
     7  package iofmt
     8  
     9  import (
    10  	"bytes"
    11  	"io"
    12  )
    13  
    14  var newline = []byte{'\n'}
    15  
    16  // prefixWriter is an io.Writer that outputs a prefix before each line.
    17  type prefixWriter struct {
    18  	w          io.Writer
    19  	prefix     string
    20  	needPrefix bool
    21  }
    22  
    23  // PrefixWriter returns a new io.Writer that copies its writes
    24  // to the provided io.Writer, adding a prefix at the beginning
    25  // of each line.
    26  func PrefixWriter(w io.Writer, prefix string) io.Writer {
    27  	return &prefixWriter{w: w, prefix: prefix, needPrefix: true}
    28  }
    29  
    30  func (w *prefixWriter) Write(p []byte) (n int, err error) {
    31  	if w.needPrefix {
    32  		if _, err := io.WriteString(w.w, w.prefix); err != nil {
    33  			return 0, err
    34  		}
    35  		w.needPrefix = false
    36  	}
    37  	for {
    38  		i := bytes.Index(p, newline)
    39  		switch i {
    40  		case len(p) - 1:
    41  			w.needPrefix = true
    42  			fallthrough
    43  		case -1:
    44  			m, err := w.w.Write(p)
    45  			return n + m, err
    46  		default:
    47  			m, err := w.w.Write(p[:i+1])
    48  			n += m
    49  			if err != nil {
    50  				return n, err
    51  			}
    52  			_, err = io.WriteString(w.w, w.prefix)
    53  			if err != nil {
    54  				return n, err
    55  			}
    56  			p = p[i+1:]
    57  		}
    58  	}
    59  }