pretty print download progress when run from console

diff --git a/youtube.cgi b/youtube.cgi
index 293db8ab74d013e7306048bdd48c6eb56aca50af..17fcf1697b517ec053abb2ff98a6dfd74202ea93 100755 (executable)
--- 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)