github.com/TrenchBoot/u-root@v6.0.1-0.20200623220532-550e36da3258+incompatible/pkg/uio/linewriter.go (about)

     1  // Copyright 2019 the u-root Authors. All rights reserved
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package uio
     6  
     7  import (
     8  	"bytes"
     9  	"io"
    10  )
    11  
    12  // LineWriter processes one line of log output at a time.
    13  type LineWriter interface {
    14  	// OneLine is always called with exactly one line of output.
    15  	OneLine(b []byte)
    16  }
    17  
    18  // FullLineWriter returns an io.Writer that waits for a full line of prints
    19  // before calling w.Write on one line each.
    20  func FullLineWriter(w LineWriter) io.WriteCloser {
    21  	return &fullLineWriter{w: w}
    22  }
    23  
    24  type fullLineWriter struct {
    25  	w      LineWriter
    26  	buffer []byte
    27  }
    28  
    29  func (fsw *fullLineWriter) printBuf() {
    30  	bufs := bytes.Split(fsw.buffer, []byte{'\n'})
    31  	for _, buf := range bufs {
    32  		if len(buf) != 0 {
    33  			fsw.w.OneLine(buf)
    34  		}
    35  	}
    36  	fsw.buffer = nil
    37  }
    38  
    39  // Write implements io.Writer and buffers p until at least one full line is
    40  // received.
    41  func (fsw *fullLineWriter) Write(p []byte) (int, error) {
    42  	i := bytes.LastIndexByte(p, '\n')
    43  	if i == -1 {
    44  		fsw.buffer = append(fsw.buffer, p...)
    45  	} else {
    46  		fsw.buffer = append(fsw.buffer, p[:i]...)
    47  		fsw.printBuf()
    48  		fsw.buffer = append([]byte{}, p[i:]...)
    49  	}
    50  	return len(p), nil
    51  }
    52  
    53  // Close implements io.Closer and flushes the buffer.
    54  func (fsw *fullLineWriter) Close() error {
    55  	fsw.printBuf()
    56  	return nil
    57  }