github.com/hasnat/dolt/go@v0.0.0-20210628190320-9eb5d843fbb7/store/util/writers/prefix_writer.go (about)

     1  // Copyright 2019 Dolthub, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  //
    15  // This file incorporates work covered by the following copyright and
    16  // permission notice:
    17  //
    18  // Copyright 2016 Attic Labs, Inc. All rights reserved.
    19  // Licensed under the Apache License, version 2.0:
    20  // http://www.apache.org/licenses/LICENSE-2.0
    21  
    22  package writers
    23  
    24  import "io"
    25  
    26  // PrefixWriter makes it easy to prefix lines with a custom prefix. Each time
    27  // it writes a byte after a newline('\n') character it calls PrefixFunc() to get
    28  // the byte slice that should be written. |NeedsPrefix| can be set to true to
    29  // cause a prefix to be written immediately. This is useful for causing a prefix
    30  // to get written on the first line.
    31  type PrefixWriter struct {
    32  	Dest        io.Writer
    33  	PrefixFunc  func(w *PrefixWriter) []byte
    34  	NeedsPrefix bool
    35  	NumLines    uint32
    36  }
    37  
    38  // Write() will add a prefix to the beginning of each line. It obtains the
    39  // prefix by call |PrefixFunc(w *PrefixWriter)| before printing out any character
    40  // following a newLine. Callers can force a prefix to be printed out before the
    41  // first character in |data| by setting NeedsPrefix to true. Conversely, callers
    42  // can suppress prefixes from being printed by setting NeedsPrefix to false.
    43  
    44  func (w *PrefixWriter) Write(data []byte) (int, error) {
    45  	writtenCnt := 0
    46  	for i, b := range data {
    47  		if w.NeedsPrefix {
    48  			w.NeedsPrefix = false
    49  			d1 := w.PrefixFunc(w)
    50  			cnt, err := w.Dest.Write(d1)
    51  			writtenCnt += cnt
    52  			if err != nil {
    53  				return writtenCnt, err
    54  			}
    55  		}
    56  		if b == byte('\n') {
    57  			w.NumLines++
    58  			w.NeedsPrefix = true
    59  		}
    60  		cnt, err := w.Dest.Write(data[i : i+1])
    61  		writtenCnt += cnt
    62  		if err != nil {
    63  			return writtenCnt, err
    64  		}
    65  	}
    66  	return writtenCnt, nil
    67  }