hentai/downloader.py

   1 #coding: utf-8
   2 import os
   3 import sys
   4 import threading
   5 import Queue
   6 import requests
   7 from urlparse import urlparse
   8 from hentai.logger import logger
   9
  10
  11 class Downloader(object):
  12     _instance = None
  13
  14     def __new__(cls, *args, **kwargs):
  15         if not cls._instance:
  16             cls._instance = super(Downloader, cls).__new__(cls, *args, **kwargs)
  17         return cls._instance
  18
  19     def __init__(self, path='', thread=1):
  20         if not isinstance(thread, (int, )) or thread < 1 or thread > 10:
  21             raise ValueError('Invalid threads count')
  22         self.path = str(path)
  23         self.thread_count = thread
  24         self.threads = []
  25
  26     def _download(self, url, folder='', filename=''):
  27         if not os.path.exists(folder):
  28             try:
  29                 os.mkdir(folder)
  30             except os.error, e:
  31                 logger.error('Error %s' % str(e))
  32                 sys.exit()
  33
  34         filename = filename if filename else os.path.basename(urlparse(url).path)
  35         try:
  36             with open(os.path.join(folder, filename), "wb") as f:
  37                 response = requests.get(url, stream=True)
  38                 length = response.headers.get('content-length')
  39                 if length is None:
  40                     f.write(response.content)
  41                 else:
  42                     for chunk in response.iter_content(2048):
  43                         f.write(chunk)
  44         except (os.error, IOError), e:
  45             logger.error('Error %s' % e)
  46             sys.exit()
  47
  48         except Exception, e:
  49             raise e
  50
  51         logger.info('%s %s downloaded.' % (threading.currentThread().getName(), url))
  52
  53     def _download_thread(self, queue, folder=''):
  54         while True:
  55             if queue.empty():
  56                 queue.task_done()
  57                 break
  58             try:
  59                 url = queue.get(False)
  60                 logger.info('%s downloading: %s ...' % (threading.currentThread().getName(), url))
  61                 self._download(url, folder)
  62             except Queue.Empty:
  63                 break
  64
  65     def download(self, queue, folder=''):
  66         if not isinstance(folder, (str, unicode)):
  67             folder = str(folder)
  68
  69         if self.path:
  70             folder = '%s/%s' % (self.path, folder)
  71
  72         if os.path.exists(path=folder):
  73             logger.warn('Path \'%s\' already exist' % folder)
  74         else:
  75             logger.warn('Path \'%s\' not exist' % folder)
  76
  77         for i in range(self.thread_count):
  78             _ = threading.Thread(target=self._download_thread, args=(queue, folder, ))
  79             _.setDaemon(True)
  80             self.threads.append(_)
  81
  82         for thread in self.threads:
  83             thread.start()
  84
  85         for thread in self.threads:
  86             thread.join()
  87
  88         # clean threads list
  89         self.threads = []
  90         logger.log(15, u'🍺 All done, saved to \'%s\'!' % folder)