From cd683cfbac0adef929f67086ba318bde32dbda76 Mon Sep 17 00:00:00 2001
From: ricterz <ricterzheng@gmail.com>
Date: Tue, 12 May 2015 22:50:29 +0800
Subject: [PATCH] retried when fail at the first time

---
 nhentai/command.py    | 10 ++++++--
 nhentai/downloader.py | 53 ++++++++++++++++++-------------------------
 2 files changed, 30 insertions(+), 33 deletions(-)

diff --git a/nhentai/command.py b/nhentai/command.py
index 88d19c7..3bef919 100644
--- a/nhentai/command.py
+++ b/nhentai/command.py
@@ -1,6 +1,6 @@
 #!/usr/bin/env python2.7
 #coding: utf-8
-
+import signal
 from cmdline import cmd_parser, banner
 from parser import dojinshi_parser, search_parser, print_dojinshi
 from dojinshi import Dojinshi
@@ -35,7 +35,8 @@ def main():
         raise SystemExit
 
     if options.is_download:
-        downloader = Downloader(path=options.saved_path, thread=options.threads)
+        downloader = Downloader(path=options.saved_path,
+                                thread=options.threads, timeout=options.timeout)
         for dojinshi in dojinshi_list:
             dojinshi.downloader = downloader
             dojinshi.download()
@@ -45,5 +46,10 @@ def main():
     logger.log(15, u'🍺 All done.')
 
 
+def signal_handler(signal, frame):
+    logger.error('Ctrl-C signal received. Quit.')
+    raise SystemExit
+
+signal.signal(signal.SIGINT, signal_handler)
 if __name__ == '__main__':
     main()
\ No newline at end of file
diff --git a/nhentai/downloader.py b/nhentai/downloader.py
index 41f2678..fc054f6 100644
--- a/nhentai/downloader.py
+++ b/nhentai/downloader.py
@@ -1,50 +1,33 @@
 #coding: utf-8
 import os
-import sys
-import socket
-import threading
-import Queue
 import requests
 import threadpool
 from urlparse import urlparse
 from logger import logger
 
 
-# global timeout
-timeout = 30
-THREAD_TIMEOUT = 99999
-socket.setdefaulttimeout(timeout)
-
-
 class Downloader(object):
     _instance = None
-    kill_received = False
 
     def __new__(cls, *args, **kwargs):
         if not cls._instance:
             cls._instance = super(Downloader, cls).__new__(cls, *args, **kwargs)
         return cls._instance
 
-    def __init__(self, path='', thread=1):
+    def __init__(self, path='', thread=1, timeout=30):
         if not isinstance(thread, (int, )) or thread < 1 or thread > 10:
             raise ValueError('Invalid threads count')
         self.path = str(path)
         self.thread_count = thread
         self.threads = []
+        self.timeout = timeout
 
-    def _download(self, url, folder='', filename=''):
-        if not os.path.exists(folder):
-            try:
-                os.mkdir(folder)
-            except os.error, e:
-                logger.critical('Error: %s' % str(e))
-                sys.exit()
-
+    def _download(self, url, folder='', filename='', retried=False):
         logger.info('Start downloading: %s ...' % url)
         filename = filename if filename else os.path.basename(urlparse(url).path)
         try:
             with open(os.path.join(folder, filename), "wb") as f:
-                response = requests.get(url, stream=True, timeout=timeout)
+                response = requests.get(url, stream=True, timeout=self.timeout)
                 length = response.headers.get('content-length')
                 if length is None:
                     f.write(response.content)
@@ -52,13 +35,20 @@ class Downloader(object):
                     for chunk in response.iter_content(2048):
                         f.write(chunk)
         except (os.error, IOError), e:
-            logger.critical('Error: %s' % str(e))
-            sys.exit()
+            if not retried:
+                logger.error('Error: %s, retrying' % str(e))
+                return self._download(url=url, folder=folder, filename=filename, retried=True)
+            else:
+                return None
         except Exception, e:
+            logger.critical('CRITICAL: %s' % str(e))
             raise e
         return url
 
     def _download_callback(self, request, result):
+        if not result:
+            logger.critical('Too many network errors occurred, please check your connection.')
+            raise SystemExit
         logger.log(15, '%s download successfully' % result)
 
     def download(self, queue, folder=''):
@@ -68,10 +58,15 @@ class Downloader(object):
         if self.path:
             folder = os.path.join(self.path, folder)
 
-        if os.path.exists(path=folder):
-            logger.warn('Path \'%s\' already exist' % folder)
+        if not os.path.exists(folder):
+            logger.warn('Path \'%s\' not exist.' % folder)
+            try:
+                os.mkdir(folder)
+            except os.error, e:
+                logger.critical('Error: %s' % str(e))
+                raise SystemExit
         else:
-            logger.warn('Path \'%s\' not exist' % folder)
+            logger.warn('Path \'%s\' already exist.' % folder)
 
         queue = [([url], {'folder': folder}) for url in queue]
 
@@ -79,8 +74,4 @@ class Downloader(object):
         requests_ = threadpool.makeRequests(self._download, queue, self._download_callback)
         [self.thread_pool.putRequest(req) for req in requests_]
 
-        try:
-            self.thread_pool.wait()
-        except KeyboardInterrupt:
-            print
-            logger.error('Ctrl-C pressed, exiting threads ...')
+        self.thread_pool.wait()