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