jben/jben/download_thread.py

   1 # -*- coding: utf-8 -*-
   2
   3 from __future__ import with_statement
   4
   5 import threading, Queue, urllib2
   6 import sys
   7
   8
   9 class DownloadThread(threading.Thread):
  10
  11     """Thread object for downloading files.
  12
  13     Takes a url string and output file name as required parameters.
  14
  15     """
  16
  17     CONNECTING, DOWNLOADING, DONE, ERROR = range(4)
  18     ABORT = 1
  19
  20     def __init__(self, url, fname, chunk_size=16384, timeout=None):
  21         threading.Thread.__init__(self)
  22         self.out_queue = Queue.Queue(0)
  23         self.in_queue = Queue.Queue(0)
  24         self.url = url
  25         self.fname = fname
  26         self.chunk_size = chunk_size
  27         self.timeout = timeout
  28
  29     def run(self):
  30         try:
  31             with open(self.fname, "wb") as ofile:
  32                 progress = 0
  33                 self.out_queue.put((self.CONNECTING, 0))
  34                 if not (self.timeout is None):
  35                     resp = urllib2.urlopen(self.url, timeout=self.timeout)
  36                 else:
  37                     resp = urllib2.urlopen(self.url)
  38                 while True:
  39                     try:
  40                         event = self.in_queue.get(block=False)
  41                     except Queue.Empty:
  42                         event = None
  43                     if event == self.ABORT:
  44                         raise Exception("Aborted by parent thread")
  45                     self.out_queue.put((self.DOWNLOADING, progress))
  46                     d = resp.read(self.chunk_size)
  47                     if not d:
  48                         break
  49                     ofile.write(d)
  50                     progress += len(d)
  51                 self.out_queue.put((self.DONE, progress))
  52         except Exception, e:
  53             self.out_queue.put((self.ERROR, e))
  54
  55
  56 def main():
  57     assert len(sys.argv) >= 2, "Please specify a URL to download."
  58     assert len(sys.argv) >= 3, "Please specify a file name to save to."
  59     dt = DownloadThread(sys.argv[1], sys.argv[2], timeout=5)
  60     print "Starting thread..."
  61     dt.start()
  62     try:
  63         while True:
  64             (status, progress) = dt.out_queue.get(block=True)
  65             if status == dt.CONNECTING:
  66                 print "Connecting..."
  67             elif status == dt.DOWNLOADING:
  68                 print "Downloading: %d bytes received" % progress
  69             elif status == dt.DONE:
  70                 break
  71             elif status == dt.ERROR:
  72                 # progress is an exception object
  73                 raise progress
  74     except KeyboardInterrupt, e:
  75         print "Keyboard interrupt received; aborting."
  76         dt.in_queue.put(dt.ABORT)
  77     dt.join()
  78     print "Done!  %d bytes received!" % progress
  79
  80 if __name__ == "__main__":
  81     main()