| // Copyright 2015 The Vanadium Authors. All rights reserved. |
| // Use of this source code is governed by a BSD-style |
| // license that can be found in the LICENSE file. |
| |
| package textutil |
| |
| import ( |
| "bytes" |
| "io" |
| "unicode/utf8" |
| ) |
| |
| // WriteFlusher is the interface that groups the basic Write and Flush methods. |
| // |
| // Flush is typically provided when Write calls perform buffering; Flush |
| // immediately outputs the buffered data. Flush must be called after the last |
| // call to Write, and may be called an arbitrary number of times before the last |
| // Write. |
| // |
| // If the type is also a Closer, Close implies a Flush call. |
| type WriteFlusher interface { |
| io.Writer |
| Flush() error |
| } |
| |
| // WriteFlushCloser is the interface that groups the basic Write, Flush and |
| // Close methods. |
| type WriteFlushCloser interface { |
| WriteFlusher |
| io.Closer |
| } |
| |
| // PrefixWriter returns an io.Writer that wraps w, where the prefix is written |
| // out immediately before the first non-empty Write call. |
| func PrefixWriter(w io.Writer, prefix string) io.Writer { |
| return &prefixWriter{w, []byte(prefix)} |
| } |
| |
| type prefixWriter struct { |
| w io.Writer |
| prefix []byte |
| } |
| |
| func (w *prefixWriter) Write(data []byte) (int, error) { |
| if w.prefix != nil && len(data) > 0 { |
| w.w.Write(w.prefix) |
| w.prefix = nil |
| } |
| return w.w.Write(data) |
| } |
| |
| // PrefixLineWriter returns a WriteFlushCloser that wraps w. Any occurrence of |
| // EOL (\f, \n, \r, \v, LineSeparator or ParagraphSeparator) causes the |
| // preceeding line to be written to w, with the given prefix. Data without EOL |
| // is buffered until the next EOL, or Flush or Close call. A single Write call |
| // may result in zero or more Write calls on the underlying writer. |
| func PrefixLineWriter(w io.Writer, prefix string) WriteFlushCloser { |
| return &prefixLineWriter{w, []byte(prefix), nil} |
| } |
| |
| type prefixLineWriter struct { |
| w io.Writer |
| prefix []byte |
| buf []byte |
| } |
| |
| const eolRunesAsString = "\f\n\r\v" + string(LineSeparator) + string(ParagraphSeparator) |
| |
| func (w *prefixLineWriter) Write(data []byte) (int, error) { |
| totalLen := len(data) |
| for len(data) > 0 { |
| index := bytes.IndexAny(data, eolRunesAsString) |
| if index == -1 { |
| // No EOL: buffer remaining data. |
| // TODO(toddw): Flush at a max size, to avoid unbounded growth? |
| w.buf = append(w.buf, data...) |
| return totalLen, nil |
| } |
| // Saw EOL: write prefix, buffer, and data including EOL. |
| if _, err := w.w.Write(w.prefix); err != nil { |
| return totalLen - len(data), err |
| } |
| if _, err := w.w.Write(w.buf); err != nil { |
| return totalLen - len(data), err |
| } |
| w.buf = w.buf[:0] |
| _, eolSize := utf8.DecodeRune(data[index:]) |
| n, err := w.w.Write(data[:index+eolSize]) |
| data = data[n:] |
| if err != nil { |
| return totalLen - len(data), err |
| } |
| } |
| return totalLen, nil |
| } |
| |
| func (w *prefixLineWriter) Flush() error { |
| if len(w.buf) > 0 { |
| if _, err := w.w.Write(w.prefix); err != nil { |
| return err |
| } |
| if _, err := w.w.Write(w.buf); err != nil { |
| return err |
| } |
| w.buf = w.buf[:0] |
| } |
| if f, ok := w.w.(WriteFlusher); ok { |
| return f.Flush() |
| } |
| return nil |
| } |
| |
| func (w *prefixLineWriter) Close() error { |
| firstErr := w.Flush() |
| if c, ok := w.w.(io.Closer); ok { |
| if err := c.Close(); firstErr == nil { |
| firstErr = err |
| } |
| } |
| return firstErr |
| } |
| |
| // ByteReplaceWriter returns an io.Writer that wraps w, where all occurrences of |
| // the old byte are replaced with the new string on Write calls. |
| func ByteReplaceWriter(w io.Writer, old byte, new string) io.Writer { |
| return &byteReplaceWriter{w, []byte{old}, []byte(new)} |
| } |
| |
| type byteReplaceWriter struct { |
| w io.Writer |
| old, new []byte |
| } |
| |
| func (w *byteReplaceWriter) Write(data []byte) (int, error) { |
| replaced := bytes.Replace(data, w.old, w.new, -1) |
| if len(replaced) == 0 { |
| return len(data), nil |
| } |
| // Write the replaced data, and return the number of bytes in data that were |
| // written out, based on the proportion of replaced data written. The |
| // important boundary cases are: |
| // If all replaced data was written, we return n=len(data). |
| // If not all replaced data was written, we return n<len(data). |
| n, err := w.w.Write(replaced) |
| return n * len(data) / len(replaced), err |
| } |
| |
| // TODO(toddw): Add ReplaceWriter, which performs arbitrary string replacements. |
| // This will need to buffer data and have an extra Flush() method, since the old |
| // string may match across successive Write calls. |