From 9549c5f5a23dd2b78d9d83ae8ae9945e5993be57 Mon Sep 17 00:00:00 2001
From: Ricter Z <ricterzheng@gmail.com>
Date: Wed, 23 Nov 2016 22:35:56 +0800
Subject: [PATCH] fix bug

---
 nhentai/cmdline.py    | 10 +++++-----
 nhentai/downloader.py | 41 ++++++++++++++++++++++++++++-------------
 2 files changed, 33 insertions(+), 18 deletions(-)

diff --git a/nhentai/cmdline.py b/nhentai/cmdline.py
index 624b9ab..eca33ac 100644
--- a/nhentai/cmdline.py
+++ b/nhentai/cmdline.py
@@ -30,10 +30,10 @@ def banner():
 
 
 def cmd_parser():
-    parser = OptionParser('\nnhentai --search [keyword] --download'
-                          '\nNHENTAI=http://h.loli.club nhentai --id [ID ...]'
+    parser = OptionParser('\n  nhentai --search [keyword] --download'
+                          '\n  NHENTAI=http://h.loli.club nhentai --id [ID ...]'
                           '\n\nEnvironment Variable:\n'
-                          'NHENTAI                 nhentai mirror url')
+                          '  NHENTAI                 nhentai mirror url')
     parser.add_option('--download', dest='is_download', action='store_true', help='download doujinshi (for search result)')
     parser.add_option('--show-info', dest='is_show', action='store_true', help='just show the doujinshi information')
     parser.add_option('--id', type='string', dest='id', action='store', help='doujinshi ids set, e.g. 1,2,3')
@@ -79,8 +79,8 @@ def cmd_parser():
     if args.threads <= 0:
         args.threads = 1
 
-    elif args.threads > 10:
-        logger.critical('Maximum number of used threads is 10')
+    elif args.threads > 15:
+        logger.critical('Maximum number of used threads is 15')
         exit(0)
 
     if args.proxy:
diff --git a/nhentai/downloader.py b/nhentai/downloader.py
index 737a409..276e1cb 100644
--- a/nhentai/downloader.py
+++ b/nhentai/downloader.py
@@ -14,17 +14,21 @@ from nhentai.parser import request
 from nhentai.utils import Singleton
 
 
+class NhentaiImageNotExistException(Exception):
+    pass
+
+
 class Downloader(Singleton):
 
     def __init__(self, path='', thread=1, timeout=30):
-        if not isinstance(thread, (int, )) or thread < 1 or thread > 10:
+        if not isinstance(thread, (int, )) or thread < 1 or thread > 15:
             raise ValueError('Invalid threads count')
         self.path = str(path)
         self.thread_count = thread
         self.threads = []
         self.timeout = timeout
 
-    def _download(self, url, folder='', filename='', retried=False):
+    def _download(self, url, folder='', filename='', retried=0):
         logger.info('Start downloading: {0} ...'.format(url))
         filename = filename if filename else os.path.basename(urlparse(url).path)
         base_filename, extension = os.path.splitext(filename)
@@ -32,29 +36,40 @@ class Downloader(Singleton):
             with open(os.path.join(folder, base_filename.zfill(3) + extension), "wb") as f:
                 response = request('get', url, stream=True, timeout=self.timeout)
                 if response.status_code != 200:
-                    logger.warning('Warning: url: {} return status code 404'.format(url))
+                    raise NhentaiImageNotExistException
                 length = response.headers.get('content-length')
                 if length is None:
                     f.write(response.content)
                 else:
                     for chunk in response.iter_content(2048):
                         f.write(chunk)
+
         except requests.HTTPError as e:
-            if not retried:
-                logger.error('Error: {0}, retrying'.format(str(e)))
-                return self._download(url=url, folder=folder, filename=filename, retried=True)
+            if retried < 3:
+                logger.warning('Warning: {0}, retrying({1}) ...'.format(str(e), retried))
+                return 0, self._download(url=url, folder=folder, filename=filename, retried=retried+1)
             else:
-                return None
+                return 0, None
+
+        except NhentaiImageNotExistException as e:
+            os.remove(os.path.join(folder, base_filename.zfill(3) + extension))
+            return -1, url
+
         except Exception as e:
             logger.critical(str(e))
-            return None
-        return url
+            return 0, None
+
+        return 1, url
 
     def _download_callback(self, request, result):
-        if not result:
-            logger.critical('Too many errors occurred, quit.')
+        result, data = result
+        if result == 0:
+            logger.critical('fatal errors occurred, quit.')
             exit(1)
-        logger.log(15, '{0} download successfully'.format(result))
+        elif result == -1:
+            logger.warning('url {} return status code 404'.format(data))
+        else:
+            logger.log(15, '{0} download successfully'.format(data))
 
     def download(self, queue, folder=''):
         if not isinstance(folder, (text)):
@@ -68,7 +83,7 @@ class Downloader(Singleton):
             try:
                 os.makedirs(folder)
             except EnvironmentError as e:
-                logger.critical('Error: {0}'.format(str(e)))
+                logger.critical('{0}'.format(str(e)))
                 exit(1)
         else:
             logger.warn('Path \'{0}\' already exist.'.format(folder))