Codebase list golang-github-vbauerster-mpb / 5e71dad
refactoring: BarClearOnComplete and BarExtender Vladimir Bauer 6 years ago
4 changed file(s) with 112 addition(s) and 93 deletion(s). Raw diff Collapse all Expand all
+52
-55
bar.go less more
6060 recoveredPanic interface{}
6161 }
6262
63 type extFunc func(in io.Reader, tw int, st *decor.Statistics) (out io.Reader, lines int)
64
6365 type bState struct {
6466 filler Filler
65 extender Filler
6667 id int
6768 width int
6869 total int64
7071 trimSpace bool
7172 toComplete bool
7273 completeFlushed bool
73 noBufBOnComplete bool
7474 noPop bool
7575 aDecorators []decor.Decorator
7676 pDecorators []decor.Decorator
7979 shutdownListeners []decor.ShutdownListener
8080 averageAdjusters []decor.AverageAdjuster
8181 bufP, bufB, bufA *bytes.Buffer
82 bufE *bytes.Buffer
82 extender extFunc
8383
8484 // priority overrides *Bar's priority, if set
8585 priority int
9292 }
9393
9494 func newBar(container *Progress, bs *bState) *Bar {
95
96 bs.bufP = bytes.NewBuffer(make([]byte, 0, bs.width))
97 bs.bufB = bytes.NewBuffer(make([]byte, 0, bs.width))
98 bs.bufA = bytes.NewBuffer(make([]byte, 0, bs.width))
99 if bs.extender != nil {
100 bs.bufE = bytes.NewBuffer(make([]byte, 0, bs.width))
101 }
102
10395 logPrefix := fmt.Sprintf("%sbar#%02d ", container.dlogger.Prefix(), bs.id)
10496 ctx, cancel := context.WithCancel(container.ctx)
97
10598 bar := &Bar{
10699 container: container,
107100 priority: bs.priority,
298291 return <-b.completed
299292 case <-b.done:
300293 return true
301 }
302 }
303
304 func (b *Bar) wSyncTable() [][]chan int {
305 select {
306 case b.operateState <- func(s *bState) { b.syncTableCh <- s.wSyncTable() }:
307 return <-b.syncTableCh
308 case <-b.done:
309 return b.cacheState.wSyncTable()
310294 }
311295 }
312296
346330 }
347331 }()
348332
349 frame := s.draw(tw)
350
351 if s.extender != nil {
352 s.extender.Fill(s.bufE, tw, newStatistics(s))
353 b.extendedLines = countLines(s.bufE.Bytes())
354 frame = io.MultiReader(frame, s.bufE)
355 }
333 st := newStatistics(s)
334 frame := s.draw(tw, st)
335 frame, b.extendedLines = s.extender(frame, tw, st)
356336
357337 b.toShutdown = s.toComplete && !s.completeFlushed
358338 s.completeFlushed = s.toComplete
359
360339 b.frameCh <- frame
361340 }:
362341 case <-b.done:
363342 s := b.cacheState
364 frame := s.draw(tw)
365 if s.extender != nil {
366 s.extender.Fill(s.bufE, tw, newStatistics(s))
367 b.extendedLines = countLines(s.bufE.Bytes())
368 frame = io.MultiReader(frame, s.bufE)
369 }
343 st := newStatistics(s)
344 frame := s.draw(tw, st)
345 frame, b.extendedLines = s.extender(frame, tw, st)
370346 b.frameCh <- frame
371347 }
372348 }
375351 return strings.NewReader(fmt.Sprintf(fmt.Sprintf("%%.%dv\n", termWidth), b.recoveredPanic))
376352 }
377353
378 func (s *bState) draw(termWidth int) io.Reader {
379
380 stat := newStatistics(s)
381
354 func (b *Bar) subscribeDecorators() {
355 var amountReceivers []decor.AmountReceiver
356 var shutdownListeners []decor.ShutdownListener
357 var averageAdjusters []decor.AverageAdjuster
358 b.TraverseDecorators(func(d decor.Decorator) {
359 if d, ok := d.(decor.AmountReceiver); ok {
360 amountReceivers = append(amountReceivers, d)
361 }
362 if d, ok := d.(decor.ShutdownListener); ok {
363 shutdownListeners = append(shutdownListeners, d)
364 }
365 if d, ok := d.(decor.AverageAdjuster); ok {
366 averageAdjusters = append(averageAdjusters, d)
367 }
368 })
369 b.operateState <- func(s *bState) {
370 s.amountReceivers = amountReceivers
371 s.shutdownListeners = shutdownListeners
372 s.averageAdjusters = averageAdjusters
373 }
374 }
375
376 func (b *Bar) refreshNowTillShutdown() {
377 for {
378 select {
379 case b.container.forceRefresh <- time.Now():
380 case <-b.done:
381 return
382 }
383 }
384 }
385
386 func (b *Bar) wSyncTable() [][]chan int {
387 select {
388 case b.operateState <- func(s *bState) { b.syncTableCh <- s.wSyncTable() }:
389 return <-b.syncTableCh
390 case <-b.done:
391 return b.cacheState.wSyncTable()
392 }
393 }
394
395 func (s *bState) draw(termWidth int, stat *decor.Statistics) io.Reader {
382396 for _, d := range s.pDecorators {
383397 s.bufP.WriteString(d.Decor(stat))
384398 }
388402 }
389403
390404 s.bufA.WriteByte('\n')
391 if s.noBufBOnComplete && s.completeFlushed {
392 return io.MultiReader(s.bufP, s.bufA)
393 }
394405
395406 prependCount := utf8.RuneCount(s.bufP.Bytes())
396407 appendCount := utf8.RuneCount(s.bufA.Bytes()) - 1
433444 return table
434445 }
435446
436 func (b *Bar) refreshNowTillShutdown() {
437 for {
438 select {
439 case b.container.forceRefresh <- time.Now():
440 case <-b.done:
441 return
442 }
443 }
444 }
445
446447 func newStatistics(s *bState) *decor.Statistics {
447448 return &decor.Statistics{
448449 ID: s.id,
451452 Current: s.current,
452453 }
453454 }
454
455 func countLines(b []byte) int {
456 return bytes.Count(b, []byte("\n"))
457 }
+28
-3
bar_option.go less more
00 package mpb
11
22 import (
3 "bytes"
4 "io"
5
36 "github.com/vbauerster/mpb/v4/decor"
47 )
58
7780 // BarClearOnComplete clears bar part of bar line on complete event.
7881 func BarClearOnComplete() BarOption {
7982 return func(s *bState) {
80 s.noBufBOnComplete = true
83 s.filler = makeClearOnCompleteFiller(s.filler)
84 }
85 }
86
87 func makeClearOnCompleteFiller(filler Filler) FillerFunc {
88 return func(w io.Writer, width int, st *decor.Statistics) {
89 if st.Completed {
90 w.Write([]byte{})
91 } else {
92 filler.Fill(w, width, st)
93 }
8194 }
8295 }
8396
93106 // BarExtender is an option to extend bar to the next new line, with
94107 // arbitrary output.
95108 func BarExtender(extender Filler) BarOption {
96 return func(s *bState) {
97 s.extender = extender
109 if extender == nil {
110 return nil
111 }
112 return func(s *bState) {
113 s.extender = makeExtFunc(extender)
114 }
115 }
116
117 func makeExtFunc(extender Filler) extFunc {
118 buf := new(bytes.Buffer)
119 nl := []byte("\n")
120 return func(r io.Reader, tw int, st *decor.Statistics) (io.Reader, int) {
121 extender.Fill(buf, tw, st)
122 return io.MultiReader(r, buf), bytes.Count(buf.Bytes(), nl)
98123 }
99124 }
100125
+2
-2
draw_test.go less more
394394 }
395395 }
396396 tmpBuf.Reset()
397 tmpBuf.ReadFrom(s.draw(termWidth))
397 tmpBuf.ReadFrom(s.draw(termWidth, newStatistics(s)))
398398 by := tmpBuf.Bytes()
399399 by = by[:len(by)-1]
400400
412412
413413 func newTestState() *bState {
414414 s := &bState{
415 filler: newDefaultBarFiller(),
415 filler: NewBarFiller(),
416416 bufP: new(bytes.Buffer),
417417 bufB: new(bytes.Buffer),
418418 bufA: new(bytes.Buffer),
+30
-33
progress.go less more
00 package mpb
11
22 import (
3 "bytes"
34 "container/heap"
45 "context"
56 "io"
6465 // context. It's not possible to reuse instance after *Progress.Wait()
6566 // method has been called.
6667 func NewWithContext(ctx context.Context, options ...ContainerOption) *Progress {
67
6868 s := &pState{
6969 bHeap: priorityQueue{},
7070 width: pwidth,
9090 done: make(chan struct{}),
9191 dlogger: log.New(s.debugOut, "[mpb] ", log.Lshortfile),
9292 }
93
9394 p.cwg.Add(1)
9495 go p.serve(s, cwriter.New(s.output))
9596 return p
115116 result := make(chan *Bar)
116117 select {
117118 case p.operateState <- func(ps *pState) {
118 bs := &bState{
119 total: total,
120 filler: filler,
121 priority: ps.idCount,
122 id: ps.idCount,
123 width: ps.width,
124 debugOut: ps.debugOut,
125 }
126 for _, opt := range options {
127 if opt != nil {
128 opt(bs)
129 }
130 }
119 bs := ps.makeBarState(total, filler, options...)
131120 bar := newBar(p, bs)
132121 if bs.runningBar != nil {
133122 bs.runningBar.noPop = true
139128 ps.idCount++
140129 result <- bar
141130 }:
142 var amountReceivers []decor.AmountReceiver
143 var shutdownListeners []decor.ShutdownListener
144 var averageAdjusters []decor.AverageAdjuster
145131 bar := <-result
146 bar.TraverseDecorators(func(d decor.Decorator) {
147 if d, ok := d.(decor.AmountReceiver); ok {
148 amountReceivers = append(amountReceivers, d)
149 }
150 if d, ok := d.(decor.ShutdownListener); ok {
151 shutdownListeners = append(shutdownListeners, d)
152 }
153 if d, ok := d.(decor.AverageAdjuster); ok {
154 averageAdjusters = append(averageAdjusters, d)
155 }
156 })
157 bar.operateState <- func(s *bState) {
158 s.amountReceivers = amountReceivers
159 s.shutdownListeners = shutdownListeners
160 s.averageAdjusters = averageAdjusters
161 }
132 bar.subscribeDecorators()
162133 return bar
163134 case <-p.done:
164135 p.bwg.Done()
357328 }
358329 }
359330
331 func (s *pState) makeBarState(total int64, filler Filler, options ...BarOption) *bState {
332 bs := &bState{
333 total: total,
334 filler: filler,
335 priority: s.idCount,
336 id: s.idCount,
337 width: s.width,
338 debugOut: s.debugOut,
339 extender: func(r io.Reader, tw int, st *decor.Statistics) (io.Reader, int) {
340 return r, 0
341 },
342 }
343
344 for _, opt := range options {
345 if opt != nil {
346 opt(bs)
347 }
348 }
349
350 bs.bufP = bytes.NewBuffer(make([]byte, 0, bs.width))
351 bs.bufB = bytes.NewBuffer(make([]byte, 0, bs.width))
352 bs.bufA = bytes.NewBuffer(make([]byte, 0, bs.width))
353
354 return bs
355 }
356
360357 func syncWidth(matrix map[int][]chan int) {
361358 for _, column := range matrix {
362359 column := column