From 9230a1b9df1f757f6fadc27933718274c0d7114c Mon Sep 17 00:00:00 2001
From: Matt Joiner <anacrolix@gmail.com>
Date: Sun, 22 Nov 2015 18:45:06 +1100
Subject: [PATCH] Use uiprogress to display a progress bar per torrent

---
 cmd/torrent/main.go | 149 ++++++++++++++++++++------------------------
 1 file changed, 68 insertions(+), 81 deletions(-)

diff --git a/cmd/torrent/main.go b/cmd/torrent/main.go
index 0bbd0ec4..4adda93e 100644
--- a/cmd/torrent/main.go
+++ b/cmd/torrent/main.go
@@ -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"
 		}
-	}
-	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
+		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())))
 		}
-		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)
-		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))
+}
+
+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?!")
+	}
 	if opts.Seed {
 		select {}
 	}