Codebase list golang-github-vbauerster-mpb / e45b2f5
introduce cwriter package Vladimir Bauer 9 years ago
5 changed file(s) with 161 addition(s) and 2 deletion(s). Raw diff Collapse all Expand all
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 }
77 "sync"
88 "time"
99
10 "github.com/vbauerster/uilive"
10 "github.com/vbauerster/mpb/cwriter"
1111 )
1212
1313 type opType uint
116116 func (p *Progress) server() {
117117 t := time.NewTicker(refreshRate * time.Millisecond)
118118 bars := make([]*Bar, 0, 4)
119 lw := uilive.New(p.out)
119 lw := cwriter.New(p.out)
120120 for {
121121 select {
122122 case op, ok := <-p.op: