diff --git a/cwriter/writer.go b/cwriter/writer.go new file mode 100644 index 0000000..138a519 --- /dev/null +++ b/cwriter/writer.go @@ -0,0 +1,50 @@ +package cwriter + +import ( + "bytes" + "io" +) + +// ESC is the ASCII code for escape character +const ESC = 27 + +// Writer is a buffered the writer that updates the terminal. +// The contents of writer will be flushed when Flush is called. +type Writer struct { + // Out is the writer to write to + out io.Writer + + buf bytes.Buffer + lineCount int +} + +// New returns a new Writer with defaults +func New(w io.Writer) *Writer { + return &Writer{ + out: w, + } +} + +func (w *Writer) Flush() error { + // do nothing if buffer is empty + if len(w.buf.Bytes()) == 0 { + return nil + } + w.clearLines() + + lines := 0 + for _, b := range w.buf.Bytes() { + if b == '\n' { + lines++ + } + } + w.lineCount = lines + _, err := w.out.Write(w.buf.Bytes()) + w.buf.Reset() + return err +} + +// Write save the contents of b to its buffers. The only errors returned are ones encountered while writing to the underlying buffer. +func (w *Writer) Write(b []byte) (n int, err error) { + return w.buf.Write(b) +} diff --git a/cwriter/writer_posix.go b/cwriter/writer_posix.go new file mode 100644 index 0000000..854b47b --- /dev/null +++ b/cwriter/writer_posix.go @@ -0,0 +1,14 @@ +// +build !windows + +package cwriter + +import ( + "fmt" +) + +func (w *Writer) clearLines() { + for i := 0; i < w.lineCount; i++ { + fmt.Fprintf(w.out, "%c[%dA", ESC, 0) // move the cursor up + fmt.Fprintf(w.out, "%c[2K\r", ESC) // clear the line + } +} diff --git a/cwriter/writer_test.go b/cwriter/writer_test.go new file mode 100644 index 0000000..8aec781 --- /dev/null +++ b/cwriter/writer_test.go @@ -0,0 +1,20 @@ +package cwriter + +import ( + "bytes" + "fmt" + "testing" +) + +func TestWriter(t *testing.T) { + b := &bytes.Buffer{} + w := New(b) + for i := 0; i < 2; i++ { + fmt.Fprintln(w, "foo") + } + w.Flush() + want := "foo\nfoo\n" + if b.String() != want { + t.Fatalf("want %q, got %q", want, b.String()) + } +} diff --git a/cwriter/writer_windows.go b/cwriter/writer_windows.go new file mode 100644 index 0000000..ec5b40f --- /dev/null +++ b/cwriter/writer_windows.go @@ -0,0 +1,75 @@ +// +build windows + +package cwriter + +import ( + "fmt" + "syscall" + "unsafe" + + "github.com/mattn/go-isatty" +) + +var kernel32 = syscall.NewLazyDLL("kernel32.dll") + +var ( + procGetConsoleScreenBufferInfo = kernel32.NewProc("GetConsoleScreenBufferInfo") + procSetConsoleCursorPosition = kernel32.NewProc("SetConsoleCursorPosition") + procFillConsoleOutputCharacter = kernel32.NewProc("FillConsoleOutputCharacterW") + procFillConsoleOutputAttribute = kernel32.NewProc("FillConsoleOutputAttribute") +) + +type short int16 +type dword uint32 +type word uint16 + +type coord struct { + x short + y short +} + +type smallRect struct { + left short + top short + right short + bottom short +} + +type consoleScreenBufferInfo struct { + size coord + cursorPosition coord + attributes word + window smallRect + maximumWindowSize coord +} + +func (w *Writer) clearLines() { + f, ok := w.Out.(FdWriter) + if ok && !isatty.IsTerminal(f.Fd()) { + ok = false + } + if !ok { + for i := 0; i < w.lineCount; i++ { + fmt.Fprintf(w.Out, "%c[%dA", ESC, 0) // move the cursor up + fmt.Fprintf(w.Out, "%c[2K\r", ESC) // clear the line + } + return + } + fd := f.Fd() + var csbi consoleScreenBufferInfo + procGetConsoleScreenBufferInfo.Call(fd, uintptr(unsafe.Pointer(&csbi))) + + for i := 0; i < w.lineCount; i++ { + // move the cursor up + csbi.cursorPosition.y-- + procSetConsoleCursorPosition.Call(fd, uintptr(*(*int32)(unsafe.Pointer(&csbi.cursorPosition)))) + // clear the line + cursor := coord{ + x: csbi.window.left, + y: csbi.window.top + csbi.cursorPosition.y, + } + var count, w dword + count = dword(csbi.size.x) + procFillConsoleOutputCharacter.Call(fd, uintptr(' '), uintptr(count), *(*uintptr)(unsafe.Pointer(&cursor)), uintptr(unsafe.Pointer(&w))) + } +} diff --git a/progress.go b/progress.go index 18e6c19..40682a1 100644 --- a/progress.go +++ b/progress.go @@ -8,7 +8,7 @@ "sync" "time" - "github.com/vbauerster/uilive" + "github.com/vbauerster/mpb/cwriter" ) type opType uint @@ -117,7 +117,7 @@ func (p *Progress) server() { t := time.NewTicker(refreshRate * time.Millisecond) bars := make([]*Bar, 0, 4) - lw := uilive.New(p.out) + lw := cwriter.New(p.out) for { select { case op, ok := <-p.op: