src/uilive/writer.go

   1 // This is a fork of github.com/gosuri/uilive for NNCP project
   2 // * It does not buffer all the writes, but resets the buffer
   3 //   just only for single latest line. Some terminals have
   4 //   huge CPU usage if so much data (as copied files progress)
   5 //   is printed
   6 // * By default it uses stderr
   7 // * By default it uses 10ms refresh period
   8 // * defer-s are removed for less CPU usage
   9 // * By default it uses stderr
  10 // * By default it uses stderr
  11 // * By default it uses stderr
  12 // * Removed newline/bypass related code. No Windows support
  13
  14 package uilive
  15
  16 import (
  17         "bytes"
  18         "fmt"
  19         "io"
  20         "os"
  21         "sync"
  22         "time"
  23 )
  24
  25 // ESC is the ASCII code for escape character
  26 const ESC = 27
  27
  28 // RefreshInterval is the default refresh interval to update the ui
  29 var RefreshInterval = 10 * time.Millisecond
  30
  31 var overFlowHandled bool
  32
  33 var termWidth int
  34
  35 // Out is the default output writer for the Writer
  36 var Out = os.Stdout
  37
  38 // FdWriter is a writer with a file descriptor.
  39 type FdWriter interface {
  40         io.Writer
  41         Fd() uintptr
  42 }
  43
  44 // Writer is a buffered the writer that updates the terminal. The contents of writer will be flushed on a timed interval or when Flush is called.
  45 type Writer struct {
  46         // Out is the writer to write to
  47         Out io.Writer
  48
  49         // RefreshInterval is the time the UI sould refresh
  50         RefreshInterval time.Duration
  51
  52         ticker *time.Ticker
  53         tdone  chan struct{}
  54
  55         buf bytes.Buffer
  56         mtx *sync.Mutex
  57 }
  58
  59 // New returns a new Writer with defaults
  60 func New() *Writer {
  61         termWidth, _ = getTermSize()
  62         if termWidth != 0 {
  63                 overFlowHandled = true
  64         }
  65         return &Writer{
  66                 Out:             Out,
  67                 RefreshInterval: RefreshInterval,
  68                 mtx:             &sync.Mutex{},
  69         }
  70 }
  71
  72 // clear the line and move the cursor up
  73 var clear = fmt.Sprintf("%c[%dA%c[2K", ESC, 1, ESC)
  74
  75 func (w *Writer) clearLines() {
  76         fmt.Fprint(w.Out, clear)
  77 }
  78
  79 // Flush writes to the out and resets the buffer. It should be called after the last call to Write to ensure that any data buffered in the Writer is written to output.
  80 // Any incomplete escape sequence at the end is considered complete for formatting purposes.
  81 // An error is returned if the contents of the buffer cannot be written to the underlying output stream
  82 func (w *Writer) Flush() (err error) {
  83         w.mtx.Lock()
  84         // do nothing if buffer is empty
  85         if len(w.buf.Bytes()) == 0 {
  86                 w.mtx.Unlock()
  87                 return
  88         }
  89         w.clearLines()
  90         var currentLine bytes.Buffer
  91         for _, b := range w.buf.Bytes() {
  92                 if b == '\n' {
  93                         currentLine.Reset()
  94                 } else {
  95                         currentLine.Write([]byte{b})
  96                         if overFlowHandled && currentLine.Len() > termWidth {
  97                                 currentLine.Reset()
  98                         }
  99                 }
 100         }
 101         _, err = w.Out.Write(w.buf.Bytes())
 102         w.mtx.Unlock()
 103         return
 104 }
 105
 106 // Start starts the listener in a non-blocking manner
 107 func (w *Writer) Start() {
 108         w.ticker = time.NewTicker(w.RefreshInterval)
 109         w.tdone = make(chan struct{}, 0)
 110         w.Out.Write([]byte("\n"))
 111         go w.Listen()
 112 }
 113
 114 // Stop stops the listener that updates the terminal
 115 func (w *Writer) Stop() {
 116         w.Flush()
 117         close(w.tdone)
 118 }
 119
 120 // Listen listens for updates to the writer's buffer and flushes to the out provided. It blocks the runtime.
 121 func (w *Writer) Listen() {
 122         for {
 123                 select {
 124                 case <-w.ticker.C:
 125                         if w.ticker != nil {
 126                                 w.Flush()
 127                         }
 128                 case <-w.tdone:
 129                         w.mtx.Lock()
 130                         w.ticker.Stop()
 131                         w.mtx.Unlock()
 132                         return
 133                 }
 134         }
 135 }
 136
 137 // Write save the contents of buf to the writer b. The only errors returned are ones encountered while writing to the underlying buffer.
 138 func (w *Writer) Write(buf []byte) (n int, err error) {
 139         w.mtx.Lock()
 140         w.buf.Reset()
 141         n, err = w.buf.Write(buf)
 142         w.mtx.Unlock()
 143         return
 144 }