@@ -11,8 +11,15 @@ import (
1111
1212var blocks = []string {"◐" , "◓" , "◑" , "◒" }
1313
14+ // WriterFn is a function that writes a formatted string.
15+ type WriteFn func (string , ... any ) (int , error )
16+
1417// MessageWriter implements io.Writer on top of a channel of strings.
15- type MessageWriter chan string
18+ type MessageWriter struct {
19+ ch chan string
20+ end chan struct {}
21+ printf WriteFn
22+ }
1623
1724// Write implements io.Writer for the MessageWriter.
1825func (m MessageWriter ) Write (p []byte ) (n int , err error ) {
@@ -22,74 +29,91 @@ func (m MessageWriter) Write(p []byte) (n int, err error) {
2229 if len (line ) == 0 {
2330 continue
2431 }
25- m <- line
32+ m . ch <- line
2633 }
2734 return len (p ), nil
2835}
2936
37+ // Closef closes the MessageWriter after writing a message.
38+ func (m MessageWriter ) Closef (format string , args ... interface {}) {
39+ m .ch <- fmt .Sprintf (format , args ... )
40+ m .Close ()
41+ }
42+
3043// Close closes the MessageWriter inner channel.
3144func (m MessageWriter ) Close () {
32- close (m )
45+ close (m .ch )
46+ <- m .end
3347}
3448
3549// Tracef is implemeted to comply with rig log.Logger interface.
3650func (m MessageWriter ) Tracef (format string , args ... interface {}) {
37- m <- fmt .Sprintf (format , args ... )
51+ m . ch <- fmt .Sprintf (format , args ... )
3852}
3953
4054// Debugf is implemeted to comply with rig log.Logger interface.
4155func (m MessageWriter ) Debugf (format string , args ... interface {}) {
42- m <- fmt .Sprintf (format , args ... )
56+ m . ch <- fmt .Sprintf (format , args ... )
4357}
4458
4559// Infof is implemeted to comply with rig log.Logger interface.
4660func (m MessageWriter ) Infof (format string , args ... interface {}) {
47- m <- fmt .Sprintf (format , args ... )
61+ m . ch <- fmt .Sprintf (format , args ... )
4862}
4963
5064// Warnf is implemeted to comply with rig log.Logger interface.
5165func (m MessageWriter ) Warnf (format string , args ... interface {}) {
52- m <- fmt .Sprintf (format , args ... )
66+ m . ch <- fmt .Sprintf (format , args ... )
5367}
5468
5569// Errorf is implemeted to comply with rig log.Logger interface.
5670func (m MessageWriter ) Errorf (format string , args ... interface {}) {
57- m <- fmt .Sprintf (format , args ... )
71+ m . ch <- fmt .Sprintf (format , args ... )
5872}
5973
60- // Start starts the progress bar. Returns a writer the caller can use to
61- // send us messages to be printed on the screen and a channel to indicate
62- // that the progress bar has finished its work (channel is closed at end)
63- // The progress bar will keep running until the MessageWriter is closed.
64- func Start () (MessageWriter , chan struct {}) {
65- finished := make (chan struct {})
66- mwriter := make (chan string , 1024 )
67- go func () {
68- var counter int
69- var message string
70- var ticker = time .NewTicker (50 * time .Millisecond )
71- var end bool
72- for {
73- select {
74- case msg , open := <- mwriter :
75- if ! open {
76- end = true
77- } else {
78- message = msg
79- }
80- case <- ticker .C :
81- counter ++
74+ // loop keeps reading messages from the channel and printint them
75+ // using the provided WriteFn. Exits when the channel is closed.
76+ func (m MessageWriter ) loop () {
77+ var counter int
78+ var message string
79+ var ticker = time .NewTicker (50 * time .Millisecond )
80+ var end bool
81+ for {
82+ select {
83+ case msg , open := <- m .ch :
84+ if ! open {
85+ end = true
86+ } else {
87+ message = msg
8288 }
89+ case <- ticker .C :
90+ counter ++
91+ }
8392
84- pos := counter % len (blocks )
85- if ! end {
86- fmt .Printf ("\033 [K\r %s %s " , blocks [pos ], message )
87- continue
88- }
89- fmt .Printf ("\033 [K\r ✓ %s\n " , message )
90- close (finished )
91- return
93+ pos := counter % len (blocks )
94+ if ! end {
95+ _ , _ = m .printf ("\033 [K\r %s %s " , blocks [pos ], message )
96+ continue
9297 }
93- }()
94- return mwriter , finished
98+ _ , _ = m .printf ("\033 [K\r ✓ %s\n " , message )
99+ close (m .end )
100+ return
101+ }
102+ }
103+
104+ // Start starts the progress bar. Returns a writer the caller can use to
105+ // send us messages to be printed using the provided WriteFn (if nil then
106+ // we print to stdout). MessageWriter should be closed when done.
107+ // MessageWriter is closed.
108+ func Start (fn WriteFn ) MessageWriter {
109+ if fn == nil {
110+ fn = fmt .Printf
111+ }
112+ mw := MessageWriter {
113+ ch : make (chan string , 1024 ),
114+ end : make (chan struct {}),
115+ printf : fn ,
116+ }
117+ go mw .loop ()
118+ return mw
95119}
0 commit comments