Codebase list golang-github-vbauerster-mpb / f0b60c9
New accepts context for cancellation purposes Vladimir Bauer 9 years ago
3 changed file(s) with 74 addition(s) and 45 deletion(s). Raw diff Collapse all Expand all
+18
-10
bar.go less more
2525 stateReqCh chan chan state
2626 decoratorCh chan *decorator
2727 flushedCh chan struct{}
28 removeReqCh chan struct{}
2829 done chan struct{}
2930
3031 lastState state
6566 stateReqCh: make(chan chan state),
6667 decoratorCh: make(chan *decorator),
6768 flushedCh: make(chan struct{}),
69 removeReqCh: make(chan struct{}),
6870 done: make(chan struct{}),
6971 }
7072 go b.server(ctx, wg, total)
159161 b.stateReqCh <- ch
160162 state := <-ch
161163 return state.current
162 }
163
164 func (b *Bar) stop() {
165 if !b.isDone() {
166 close(b.done)
167 }
168164 }
169165
170166 // InProgress returns true, while progress is running
237233 case state.trimRightSpace = <-b.trimRightCh:
238234 case <-b.flushedCh:
239235 if completed {
240 b.lastState = state
241 b.stop()
236 b.stop(&state)
242237 return
243238 }
239 case <-b.removeReqCh:
240 b.stop(&state)
241 return
244242 case <-ctx.Done():
245 b.lastState = state
246 b.stop()
243 b.stop(&state)
247244 return
248245 }
249246 }
247 }
248
249 func (b *Bar) stop(s *state) {
250 b.lastState = *s
251 close(b.done)
250252 }
251253
252254 func (b *Bar) draw(s state, termWidth int) []byte {
326328
327329 return buf
328330 }
331
332 // func (b *Bar) closeDone() {
333 // if !b.isDone() {
334 // close(b.done)
335 // }
336 // }
329337
330338 func (b *Bar) isDone() bool {
331339 select {
+54
-33
progress.go less more
00 package mpb
11
22 import (
3 "context"
34 "errors"
45 "io"
56 "os"
3435
3536 // Progress represents the container that renders Progress bars
3637 type Progress struct {
38 // Context for canceling bars rendering
39 ctx context.Context
3740 // WaitGroup for internal rendering sync
3841 wg *sync.WaitGroup
3942
4144 width int
4245 sort SortType
4346
44 op chan *operation
47 operationCh chan *operation
4548 rrChangeReqCh chan time.Duration
4649 outChangeReqCh chan io.Writer
47 countReqCh chan chan int
48 allDone chan struct{}
50 barCountReqCh chan chan int
51 done chan struct{}
4952 }
5053
5154 type operation struct {
5457 result chan bool
5558 }
5659
57 // New returns a new progress bar with defaults
58 func New() *Progress {
60 // New creates new Progress instance, which will orchestrate bars rendering
61 // process. It acceepts context.Context, for cancellation.
62 // If you don't plan to cancel, it is safe to feed with nil
63 func New(ctx context.Context) *Progress {
64 if ctx == nil {
65 ctx = context.Background()
66 }
5967 p := &Progress{
6068 width: 70,
61 op: make(chan *operation),
69 operationCh: make(chan *operation),
6270 rrChangeReqCh: make(chan time.Duration),
6371 outChangeReqCh: make(chan io.Writer),
64 countReqCh: make(chan chan int),
65 allDone: make(chan struct{}),
72 barCountReqCh: make(chan chan int),
73 done: make(chan struct{}),
6674 wg: new(sync.WaitGroup),
75 ctx: ctx,
6776 }
6877 go p.server(cwriter.New(os.Stdout), time.NewTicker(rr*time.Millisecond))
6978 return p
8190 // SetOut sets underlying writer of progress. Default is os.Stdout
8291 // pancis, if called on stopped Progress instance, i.e after Stop()
8392 func (p *Progress) SetOut(w io.Writer) *Progress {
84 if p.isAllDone() {
93 if p.isDone() {
8594 panic(ErrCallAfterStop)
8695 }
8796 if w == nil {
94103 // RefreshRate overrides default (30ms) refreshRate value
95104 // pancis, if called on stopped Progress instance, i.e after Stop()
96105 func (p *Progress) RefreshRate(d time.Duration) *Progress {
97 if p.isAllDone() {
106 if p.isDone() {
98107 panic(ErrCallAfterStop)
99108 }
100109 p.rrChangeReqCh <- d
101110 return p
102111 }
112
113 // func (p *Progress) WithContext(ctx context.Context) *Progress {
114 // if p.BarCount() > 0 {
115 // panic("cannot apply ctx after AddBar has been called")
116 // }
117 // if ctx == nil {
118 // panic("nil context")
119 // }
120 // p.ctx = ctx
121 // return p
122 // }
103123
104124 // WithSort sorts the bars, while redering
105125 func (p *Progress) WithSort(sort SortType) *Progress {
109129
110130 // AddBar creates a new progress bar and adds to the container
111131 // pancis, if called on stopped Progress instance, i.e after Stop()
112 func (p *Progress) AddBar(total int) *Bar {
113 if p.isAllDone() {
132 func (p *Progress) AddBar(total int64) *Bar {
133 if p.isDone() {
114134 panic(ErrCallAfterStop)
115135 }
116136 result := make(chan bool)
117 bar := newBar(total, p.width, p.wg)
118 p.op <- &operation{opBarAdd, bar, result}
137 bar := newBar(p.ctx, p.wg, total, p.width)
138 p.operationCh <- &operation{opBarAdd, bar, result}
119139 if <-result {
120140 p.wg.Add(1)
121141 }
125145 // RemoveBar removes bar at any time
126146 // pancis, if called on stopped Progress instance, i.e after Stop()
127147 func (p *Progress) RemoveBar(b *Bar) bool {
128 if p.isAllDone() {
148 if p.isDone() {
129149 panic(ErrCallAfterStop)
130150 }
131151 result := make(chan bool)
132 p.op <- &operation{opBarRemove, b, result}
152 p.operationCh <- &operation{opBarRemove, b, result}
133153 return <-result
134154 }
135155
136 // BarsCount returns bars count in the container
137 // pancis, if called on stopped Progress instance, i.e after Stop()
138 func (p *Progress) BarsCount() int {
139 if p.isAllDone() {
156 // BarCount returns bars count in the container.
157 // Pancis if called on stopped Progress instance, i.e after Stop()
158 func (p *Progress) BarCount() int {
159 if p.isDone() {
140160 panic(ErrCallAfterStop)
141161 }
142162 respCh := make(chan int)
143 p.countReqCh <- respCh
163 p.barCountReqCh <- respCh
144164 return <-respCh
145165 }
146166
147167 // Stop waits for bars to finish rendering and stops the rendering goroutine
148168 func (p *Progress) Stop() {
149 if !p.isAllDone() {
150 close(p.allDone)
151 p.wg.Wait()
152 close(p.op)
169 p.wg.Wait()
170 if !p.isDone() {
171 close(p.done)
172 close(p.operationCh)
153173 }
154174 }
155175
161181 case w := <-p.outChangeReqCh:
162182 cw.Flush()
163183 cw = cwriter.New(w)
164 case op, ok := <-p.op:
184 case op, ok := <-p.operationCh:
165185 if !ok {
166186 t.Stop()
167 for _, b := range bars {
168 b.Stop()
169 }
170187 return
171188 }
172189 switch op.kind {
179196 if b == op.bar {
180197 bars = append(bars[:i], bars[i+1:]...)
181198 ok = true
182 b.Stop()
199 b.removeReqCh <- struct{}{}
183200 break
184201 }
185202 }
186203 op.result <- ok
187204 }
188 case respCh := <-p.countReqCh:
205 case respCh := <-p.barCountReqCh:
189206 respCh <- len(bars)
190207 case <-t.C:
191208 width, _ := cwriter.TerminalWidth()
209226 case d := <-p.rrChangeReqCh:
210227 t.Stop()
211228 t = time.NewTicker(d)
229 case <-p.ctx.Done():
230 t.Stop()
231 close(p.done)
232 return
212233 }
213234 }
214235 }
215236
216 func (p *Progress) isAllDone() bool {
237 func (p *Progress) isDone() bool {
217238 select {
218 case <-p.allDone:
239 case <-p.done:
219240 return true
220241 default:
221242 return false
+2
-2
progress_test.go less more
66
77 func TestAddBar(t *testing.T) {
88 var buf bytes.Buffer
9 p := New().SetWidth(60).SetOut(&buf)
9 p := New(nil).SetWidth(60).SetOut(&buf)
1010 count := p.BarsCount()
1111 if count != 0 {
1212 t.Errorf("Count want: %q, got: %q\n", 0, count)
2323 }
2424
2525 func TestRemoveBar(t *testing.T) {
26 p := New()
26 p := New(nil)
2727 b := p.AddBar(10)
2828
2929 if !p.RemoveBar(b) {