From: James Bunton <jamesbunton@delx.net.au>
Date: Wed, 19 Sep 2012 05:18:18 +0000 (+1000)
Subject: pretty print download progress when run from console
X-Git-Url: https://code.delx.au/youtube-cgi/commitdiff_plain/85e04636dc851effce8328d54a9edb8cc6f3864e

pretty print download progress when run from console
---

diff --git a/youtube.cgi b/youtube.cgi
index 293db8a..17fcf16 100755
--- a/youtube.cgi
+++ b/youtube.cgi
@@ -11,6 +11,7 @@ import resource
 import shutil
 import subprocess
 import sys
+import time
 import urllib
 import urllib2
 import urlparse
@@ -210,6 +211,36 @@ def cgimain():
 		)
 		return
 
+def copy_with_progress(total_size, infile, outfile):
+	def pp_size(size):
+		suffixes = ["", "KiB", "MiB", "GiB"]
+		for i, suffix in enumerate(suffixes):
+			if size < 1024:
+				break
+			size /= 1024
+		return "%d %s" % (size, suffix)
+
+	start_ts = time.time()
+	last_ts = 0
+	bytes_read = 0
+	while True:
+		now = time.time()
+		if now - last_ts > 0.5:
+			last_ts = now
+			sys.stdout.write("\33[2K\r")
+			sys.stdout.write("%s / %s (%s/sec)" % (
+				pp_size(bytes_read),
+				pp_size(total_size),
+				pp_size(bytes_read / (now - start_ts)),
+			))
+			sys.stdout.flush()
+
+		buf = infile.read(32768)
+		if not buf:
+			break
+		outfile.write(buf)
+		bytes_read += len(buf)
+
 def main():
 	try:
 		url = sys.argv[1]
@@ -218,10 +249,12 @@ def main():
 		sys.exit(1)
 	doc = parse_url(url)
 	video_url, filename = get_video_url(doc)
-	data = urlopen(video_url)
+	video_data = urlopen(video_url)
 	outfile = open(filename, "w")
-	shutil.copyfileobj(data, outfile)
-	data.close()
+	total_size = int(video_data.info().getheader("Content-Length"))
+	print "Downloading", filename.encode("utf-8")
+	copy_with_progress(total_size, video_data, outfile)
+	video_data.close()
 	outfile.close()
 
 
@@ -230,5 +263,9 @@ if __name__ == "__main__":
 	if os.environ.has_key("SCRIPT_NAME"):
 		cgimain()
 	else:
-		main()
+		try:
+			main()
+		except KeyboardInterrupt:
+			print "\nExiting..."
+			sys.exit(1)