Use uiprogress to display a progress bar per torrent

This commit is contained in:
Matt Joiner 2015-11-22 18:45:06 +11:00
parent 2bdb6a14da
commit 9230a1b9df
1 changed files with 68 additions and 81 deletions

View File

@ -14,6 +14,7 @@ import (
_ "github.com/anacrolix/envpprof"
"github.com/anacrolix/tagflag"
"github.com/dustin/go-humanize"
"github.com/gosuri/uiprogress"
"github.com/anacrolix/torrent"
"github.com/anacrolix/torrent/data/mmap"
@ -35,73 +36,39 @@ func resolvedPeerAddrs(ss []string) (ret []torrent.Peer, err error) {
return
}
func bytesCompleted(tc *torrent.Client) (ret int64) {
for _, t := range tc.Torrents() {
if t.Info() != nil {
ret += t.BytesCompleted()
func torrentBar(t torrent.Torrent) {
bar := uiprogress.AddBar(1)
bar.AppendCompleted()
bar.AppendFunc(func(*uiprogress.Bar) (ret string) {
select {
case <-t.GotInfo():
default:
return "getting info"
}
if t.Seeding() {
return "seeding"
} else if t.BytesCompleted() == t.Info().TotalLength() {
return "completed"
} else {
return fmt.Sprintf("downloading (%s/%s)", humanize.Bytes(uint64(t.BytesCompleted())), humanize.Bytes(uint64(t.Info().TotalLength())))
}
return
}
// Returns an estimate of the total bytes for all torrents.
func totalBytesEstimate(tc *torrent.Client) (ret int64) {
var noInfo, hadInfo int64
for _, t := range tc.Torrents() {
info := t.Info()
if info == nil {
noInfo++
continue
}
ret += info.TotalLength()
hadInfo++
}
if hadInfo != 0 {
// Treat each torrent without info as the average of those with,
// rounded up.
ret += (noInfo*ret + hadInfo - 1) / hadInfo
}
return
}
func progressLine(tc *torrent.Client) string {
return fmt.Sprintf("\033[K%s / %s\r", humanize.Bytes(uint64(bytesCompleted(tc))), humanize.Bytes(uint64(totalBytesEstimate(tc))))
}
func main() {
log.SetFlags(log.LstdFlags | log.Lshortfile)
var opts struct {
torrent.Config `name:"Client"`
Mmap bool `help:"memory-map torrent data"`
TestPeer []*net.TCPAddr `short:"p" help:"addresses of some starting peers"`
Torrent []string `type:"pos" arity:"+" help:"torrent file path or magnet uri"`
}
tagflag.Parse(&opts, tagflag.SkipBadTypes())
clientConfig := opts.Config
if opts.Mmap {
clientConfig.TorrentDataOpener = func(info *metainfo.Info) torrent.Data {
ret, err := mmap.TorrentData(info, "")
if err != nil {
log.Fatalf("error opening torrent data for %q: %s", info.Name, err)
}
return ret
}
}
torrents := opts.Torrent
if len(torrents) == 0 {
fmt.Fprintf(os.Stderr, "no torrents specified\n")
return
}
client, err := torrent.NewClient(&clientConfig)
if err != nil {
log.Fatalf("error creating client: %s", err)
}
http.HandleFunc("/", func(w http.ResponseWriter, req *http.Request) {
client.WriteStatus(w)
})
defer client.Close()
for _, arg := range torrents {
bar.PrependFunc(func(*uiprogress.Bar) string {
return t.Name()
})
go func() {
<-t.GotInfo()
bar.Total = int(t.Info().TotalLength())
for {
bc := t.BytesCompleted()
bar.Set(int(bc))
time.Sleep(time.Second)
}
}()
}
func addTorrents(client *torrent.Client) {
for _, arg := range opts.Torrent {
t := func() torrent.Torrent {
if strings.HasPrefix(arg, "magnet:") {
t, err := client.AddMagnet(arg)
@ -122,6 +89,7 @@ func main() {
return t
}
}()
torrentBar(t)
err := t.AddPeers(func() (ret []torrent.Peer) {
for _, ta := range opts.TestPeer {
ret = append(ret, torrent.Peer{
@ -139,25 +107,44 @@ func main() {
t.DownloadAll()
}()
}
done := make(chan struct{})
go func() {
defer close(done)
}
var opts struct {
torrent.Config `name:"Client"`
Mmap bool `help:"memory-map torrent data"`
TestPeer []*net.TCPAddr `short:"p" help:"addresses of some starting peers"`
Torrent []string `type:"pos" arity:"+" help:"torrent file path or magnet uri"`
}
func main() {
log.SetFlags(log.LstdFlags | log.Lshortfile)
tagflag.Parse(&opts, tagflag.SkipBadTypes())
clientConfig := opts.Config
if opts.Mmap {
clientConfig.TorrentDataOpener = func(info *metainfo.Info) torrent.Data {
ret, err := mmap.TorrentData(info, "")
if err != nil {
log.Fatalf("error opening torrent data for %q: %s", info.Name, err)
}
return ret
}
}
client, err := torrent.NewClient(&clientConfig)
if err != nil {
log.Fatalf("error creating client: %s", err)
}
defer client.Close()
http.HandleFunc("/", func(w http.ResponseWriter, req *http.Request) {
client.WriteStatus(w)
})
uiprogress.Start()
addTorrents(client)
if client.WaitAll() {
log.Print("downloaded ALL the torrents")
} else {
log.Fatal("y u no complete torrents?!")
}
}()
ticker := time.NewTicker(time.Second)
waitDone:
for {
select {
case <-done:
break waitDone
case <-ticker.C:
os.Stdout.WriteString(progressLine(client))
}
}
if opts.Seed {
select {}
}