github.com/containerd/containerd@v22.0.0-20200918172823-438c87b8e050+incompatible/pkg/progress/writer.go (about)

     1  /*
     2     Copyright The containerd Authors.
     3  
     4     Licensed under the Apache License, Version 2.0 (the "License");
     5     you may not use this file except in compliance with the License.
     6     You may obtain a copy of the License at
     7  
     8         http://www.apache.org/licenses/LICENSE-2.0
     9  
    10     Unless required by applicable law or agreed to in writing, software
    11     distributed under the License is distributed on an "AS IS" BASIS,
    12     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13     See the License for the specific language governing permissions and
    14     limitations under the License.
    15  */
    16  
    17  package progress
    18  
    19  import (
    20  	"bytes"
    21  	"fmt"
    22  	"io"
    23  	"os"
    24  	"regexp"
    25  	"strings"
    26  
    27  	"github.com/containerd/console"
    28  )
    29  
    30  var (
    31  	regexCleanLine = regexp.MustCompile("\x1b\\[[0-9]+m[\x1b]?")
    32  )
    33  
    34  // Writer buffers writes until flush, at which time the last screen is cleared
    35  // and the current buffer contents are written. This is useful for
    36  // implementing progress displays, such as those implemented in docker and
    37  // git.
    38  type Writer struct {
    39  	buf   bytes.Buffer
    40  	w     io.Writer
    41  	lines int
    42  }
    43  
    44  // NewWriter returns a writer
    45  func NewWriter(w io.Writer) *Writer {
    46  	return &Writer{
    47  		w: w,
    48  	}
    49  }
    50  
    51  // Write the provided bytes
    52  func (w *Writer) Write(p []byte) (n int, err error) {
    53  	return w.buf.Write(p)
    54  }
    55  
    56  // Flush should be called when refreshing the current display.
    57  func (w *Writer) Flush() error {
    58  	if w.buf.Len() == 0 {
    59  		return nil
    60  	}
    61  
    62  	if err := w.clearLines(); err != nil {
    63  		return err
    64  	}
    65  	w.lines = countLines(w.buf.String())
    66  
    67  	if _, err := w.w.Write(w.buf.Bytes()); err != nil {
    68  		return err
    69  	}
    70  
    71  	w.buf.Reset()
    72  	return nil
    73  }
    74  
    75  // TODO(stevvooe): The following are system specific. Break these out if we
    76  // decide to build this package further.
    77  
    78  func (w *Writer) clearLines() error {
    79  	for i := 0; i < w.lines; i++ {
    80  		if _, err := fmt.Fprintf(w.w, "\x1b[1A\x1b[2K\r"); err != nil {
    81  			return err
    82  		}
    83  	}
    84  
    85  	return nil
    86  }
    87  
    88  // countLines in the output. If a line is longer than the console width then
    89  // an extra line is added to the count for each wrapped line. If the console
    90  // width is undefined then 0 is returned so that no lines are cleared on the next
    91  // flush.
    92  func countLines(output string) int {
    93  	con, err := console.ConsoleFromFile(os.Stdin)
    94  	if err != nil {
    95  		return 0
    96  	}
    97  	ws, err := con.Size()
    98  	if err != nil {
    99  		return 0
   100  	}
   101  	width := int(ws.Width)
   102  	if width <= 0 {
   103  		return 0
   104  	}
   105  	strlines := strings.Split(output, "\n")
   106  	lines := -1
   107  	for _, line := range strlines {
   108  		lines += (len(stripLine(line))-1)/width + 1
   109  	}
   110  	return lines
   111  }
   112  
   113  func stripLine(line string) string {
   114  	return string(regexCleanLine.ReplaceAll([]byte(line), []byte{}))
   115  }