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 17package progress 18 19import ( 20 "bytes" 21 "fmt" 22 "io" 23 "os" 24 "regexp" 25 "strings" 26 27 "github.com/containerd/console" 28) 29 30var ( 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. 38type Writer struct { 39 buf bytes.Buffer 40 w io.Writer 41 lines int 42} 43 44// NewWriter returns a writer 45func NewWriter(w io.Writer) *Writer { 46 return &Writer{ 47 w: w, 48 } 49} 50 51// Write the provided bytes 52func (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. 57func (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 78func (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. 92func 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 113func stripLine(line string) string { 114 return string(regexCleanLine.ReplaceAll([]byte(line), []byte{})) 115} 116