htpack/cmd/htpacker/pack_progress.go

package main

import (
	"github.com/logrusorgru/aurora/v4"
	"github.com/vbauerster/mpb/v4"
	"github.com/vbauerster/mpb/v4/decor"
	"src.lwithers.me.uk/go/htpack/cmd/htpacker/packer"
)

const mpbBarStyle = "[██░]"

// mpbProgress returns a new progress object that keeps the user informed via
// the visual mpb library.
func mpbProgress(ftp packer.FilesToPack) *mpbProg {
	mp := new(mpbProg)
	mp.un.max = len(ftp)

	for _, f := range ftp {
		if !f.DisableCompression && !f.DisableGzip {
			mp.gzip.max++
		}
		if !f.DisableCompression && !f.DisableBrotli {
			mp.brotli.max++
		}
	}

	mp.p = mpb.New()
	mp.un.bar = mp.p.AddBar(int64(mp.un.max),
		mpb.PrependDecorators(barName("uncompressed")),
		mpb.BarStyle(mpbBarStyle),
		mpb.AppendDecorators(&mp.un))
	if mp.gzip.max > 0 {
		mp.gzip.bar = mp.p.AddBar(int64(mp.gzip.max),
			mpb.PrependDecorators(barName("gzip")),
			mpb.BarStyle(mpbBarStyle),
			mpb.AppendDecorators(&mp.gzip))
	}
	if mp.brotli.max > 0 {
		mp.brotli.bar = mp.p.AddBar(int64(mp.brotli.max),
			mpb.PrependDecorators(barName("brotli")),
			mpb.BarStyle(mpbBarStyle),
			mpb.AppendDecorators(&mp.brotli))
	}

	return mp
}

func barName(n string) decor.Decorator {
	return decor.Name(aurora.Magenta(n).String(), decor.WCSyncWidth)
}

// mpbProg is the mpb progress tracker. It has one bar for each type of
// compression, and its methods simply dispatch onto the type-specific
// bars.
type mpbProg struct {
	un, gzip, brotli mpbProg1
	p                *mpb.Progress
}

func (mp *mpbProg) Count(_ int) {
}

func (mp *mpbProg) Begin(filename, compression string) {
	switch compression {
	case "uncompressed":
		mp.un.Begin(filename)
	case "gzip":
		mp.gzip.Begin(filename)
	case "brotli":
		mp.brotli.Begin(filename)
	default:
		return
	}
}

func (mp *mpbProg) End(filename, compression string) {
	switch compression {
	case "uncompressed":
		mp.un.End(filename)
	case "gzip":
		mp.gzip.End(filename)
	case "brotli":
		mp.brotli.End(filename)
	default:
		return
	}
}

func (mp *mpbProg) Complete() {
	mp.un.Complete()
	mp.gzip.Complete()
	mp.brotli.Complete()
	mp.p.Wait()
}

// mpbProg1 is a type-specific progress bar. In addition to holding state and
// methods for updating the bar, it also implements decor.Decor.
type mpbProg1 struct {
	max  int      // number of items we expect
	done int      // number of items completed
	cur  []string // list of currently-packing filenames
	bar  *mpb.Bar

	// embedding this type lets us implement decor.Decor
	decor.WC
}

func (mp1 *mpbProg1) Decor(stat *decor.Statistics) string {
	if stat.Completed {
		return ""
	}
	switch len(mp1.cur) {
	case 0:
		return aurora.Gray(8, "(idle)").String()
	case 1:
		return aurora.Blue(mp1.cur[0]).String()
	default:
		return aurora.Sprintf(aurora.Green("%s + %d more"), aurora.Blue(mp1.cur[0]), len(mp1.cur)-1)
	}
}

func (mp1 *mpbProg1) Begin(filename string) {
	mp1.cur = append(mp1.cur, filename)
}

func (mp1 *mpbProg1) End(filename string) {
	for i, v := range mp1.cur {
		if v == filename {
			mp1.cur[i] = mp1.cur[len(mp1.cur)-1]
			mp1.cur = mp1.cur[:len(mp1.cur)-1]
			break
		}
	}
	mp1.done++
	if mp1.bar != nil {
		mp1.bar.SetCurrent(int64(mp1.done))
	}
}

func (mp1 *mpbProg1) Complete() {
	if mp1.bar != nil {
		mp1.bar.SetTotal(int64(mp1.max), true)
	}
}