introduce cwriter package
Vladimir Bauer
9 years ago
| 0 | package cwriter | |
| 1 | ||
| 2 | import ( | |
| 3 | "bytes" | |
| 4 | "io" | |
| 5 | ) | |
| 6 | ||
| 7 | // ESC is the ASCII code for escape character | |
| 8 | const ESC = 27 | |
| 9 | ||
| 10 | // Writer is a buffered the writer that updates the terminal. | |
| 11 | // The contents of writer will be flushed when Flush is called. | |
| 12 | type Writer struct { | |
| 13 | // Out is the writer to write to | |
| 14 | out io.Writer | |
| 15 | ||
| 16 | buf bytes.Buffer | |
| 17 | lineCount int | |
| 18 | } | |
| 19 | ||
| 20 | // New returns a new Writer with defaults | |
| 21 | func New(w io.Writer) *Writer { | |
| 22 | return &Writer{ | |
| 23 | out: w, | |
| 24 | } | |
| 25 | } | |
| 26 | ||
| 27 | func (w *Writer) Flush() error { | |
| 28 | // do nothing if buffer is empty | |
| 29 | if len(w.buf.Bytes()) == 0 { | |
| 30 | return nil | |
| 31 | } | |
| 32 | w.clearLines() | |
| 33 | ||
| 34 | lines := 0 | |
| 35 | for _, b := range w.buf.Bytes() { | |
| 36 | if b == '\n' { | |
| 37 | lines++ | |
| 38 | } | |
| 39 | } | |
| 40 | w.lineCount = lines | |
| 41 | _, err := w.out.Write(w.buf.Bytes()) | |
| 42 | w.buf.Reset() | |
| 43 | return err | |
| 44 | } | |
| 45 | ||
| 46 | // Write save the contents of b to its buffers. The only errors returned are ones encountered while writing to the underlying buffer. | |
| 47 | func (w *Writer) Write(b []byte) (n int, err error) { | |
| 48 | return w.buf.Write(b) | |
| 49 | } |
| 0 | // +build !windows | |
| 1 | ||
| 2 | package cwriter | |
| 3 | ||
| 4 | import ( | |
| 5 | "fmt" | |
| 6 | ) | |
| 7 | ||
| 8 | func (w *Writer) clearLines() { | |
| 9 | for i := 0; i < w.lineCount; i++ { | |
| 10 | fmt.Fprintf(w.out, "%c[%dA", ESC, 0) // move the cursor up | |
| 11 | fmt.Fprintf(w.out, "%c[2K\r", ESC) // clear the line | |
| 12 | } | |
| 13 | } |
| 0 | package cwriter | |
| 1 | ||
| 2 | import ( | |
| 3 | "bytes" | |
| 4 | "fmt" | |
| 5 | "testing" | |
| 6 | ) | |
| 7 | ||
| 8 | func TestWriter(t *testing.T) { | |
| 9 | b := &bytes.Buffer{} | |
| 10 | w := New(b) | |
| 11 | for i := 0; i < 2; i++ { | |
| 12 | fmt.Fprintln(w, "foo") | |
| 13 | } | |
| 14 | w.Flush() | |
| 15 | want := "foo\nfoo\n" | |
| 16 | if b.String() != want { | |
| 17 | t.Fatalf("want %q, got %q", want, b.String()) | |
| 18 | } | |
| 19 | } |
| 0 | // +build windows | |
| 1 | ||
| 2 | package cwriter | |
| 3 | ||
| 4 | import ( | |
| 5 | "fmt" | |
| 6 | "syscall" | |
| 7 | "unsafe" | |
| 8 | ||
| 9 | "github.com/mattn/go-isatty" | |
| 10 | ) | |
| 11 | ||
| 12 | var kernel32 = syscall.NewLazyDLL("kernel32.dll") | |
| 13 | ||
| 14 | var ( | |
| 15 | procGetConsoleScreenBufferInfo = kernel32.NewProc("GetConsoleScreenBufferInfo") | |
| 16 | procSetConsoleCursorPosition = kernel32.NewProc("SetConsoleCursorPosition") | |
| 17 | procFillConsoleOutputCharacter = kernel32.NewProc("FillConsoleOutputCharacterW") | |
| 18 | procFillConsoleOutputAttribute = kernel32.NewProc("FillConsoleOutputAttribute") | |
| 19 | ) | |
| 20 | ||
| 21 | type short int16 | |
| 22 | type dword uint32 | |
| 23 | type word uint16 | |
| 24 | ||
| 25 | type coord struct { | |
| 26 | x short | |
| 27 | y short | |
| 28 | } | |
| 29 | ||
| 30 | type smallRect struct { | |
| 31 | left short | |
| 32 | top short | |
| 33 | right short | |
| 34 | bottom short | |
| 35 | } | |
| 36 | ||
| 37 | type consoleScreenBufferInfo struct { | |
| 38 | size coord | |
| 39 | cursorPosition coord | |
| 40 | attributes word | |
| 41 | window smallRect | |
| 42 | maximumWindowSize coord | |
| 43 | } | |
| 44 | ||
| 45 | func (w *Writer) clearLines() { | |
| 46 | f, ok := w.Out.(FdWriter) | |
| 47 | if ok && !isatty.IsTerminal(f.Fd()) { | |
| 48 | ok = false | |
| 49 | } | |
| 50 | if !ok { | |
| 51 | for i := 0; i < w.lineCount; i++ { | |
| 52 | fmt.Fprintf(w.Out, "%c[%dA", ESC, 0) // move the cursor up | |
| 53 | fmt.Fprintf(w.Out, "%c[2K\r", ESC) // clear the line | |
| 54 | } | |
| 55 | return | |
| 56 | } | |
| 57 | fd := f.Fd() | |
| 58 | var csbi consoleScreenBufferInfo | |
| 59 | procGetConsoleScreenBufferInfo.Call(fd, uintptr(unsafe.Pointer(&csbi))) | |
| 60 | ||
| 61 | for i := 0; i < w.lineCount; i++ { | |
| 62 | // move the cursor up | |
| 63 | csbi.cursorPosition.y-- | |
| 64 | procSetConsoleCursorPosition.Call(fd, uintptr(*(*int32)(unsafe.Pointer(&csbi.cursorPosition)))) | |
| 65 | // clear the line | |
| 66 | cursor := coord{ | |
| 67 | x: csbi.window.left, | |
| 68 | y: csbi.window.top + csbi.cursorPosition.y, | |
| 69 | } | |
| 70 | var count, w dword | |
| 71 | count = dword(csbi.size.x) | |
| 72 | procFillConsoleOutputCharacter.Call(fd, uintptr(' '), uintptr(count), *(*uintptr)(unsafe.Pointer(&cursor)), uintptr(unsafe.Pointer(&w))) | |
| 73 | } | |
| 74 | } |
| 7 | 7 | "sync" |
| 8 | 8 | "time" |
| 9 | 9 | |
| 10 | "github.com/vbauerster/uilive" | |
| 10 | "github.com/vbauerster/mpb/cwriter" | |
| 11 | 11 | ) |
| 12 | 12 | |
| 13 | 13 | type opType uint |
| 116 | 116 | func (p *Progress) server() { |
| 117 | 117 | t := time.NewTicker(refreshRate * time.Millisecond) |
| 118 | 118 | bars := make([]*Bar, 0, 4) |
| 119 | lw := uilive.New(p.out) | |
| 119 | lw := cwriter.New(p.out) | |
| 120 | 120 | for { |
| 121 | 121 | select { |
| 122 | 122 | case op, ok := <-p.op: |