diff --git a/nhentai/command.py b/nhentai/command.py index 7d7fbf0..e8b894f 100644 --- a/nhentai/command.py +++ b/nhentai/command.py @@ -11,7 +11,7 @@ from nhentai.doujinshi import Doujinshi from nhentai.downloader import Downloader from nhentai.logger import logger from nhentai.constant import BASE_URL -from nhentai.utils import generate_html, generate_cbz, generate_main_html, check_cookie +from nhentai.utils import generate_html, generate_cbz, generate_main_html, check_cookie, signal_handler def main(): @@ -83,12 +83,8 @@ def main(): [doujinshi.show() for doujinshi in doujinshi_list] -def signal_handler(signal, frame): - logger.error('Ctrl-C signal received. Stopping...') - exit(1) - - signal.signal(signal.SIGINT, signal_handler) + if __name__ == '__main__': main() diff --git a/nhentai/downloader.py b/nhentai/downloader.py index 81037f8..6d40016 100644 --- a/nhentai/downloader.py +++ b/nhentai/downloader.py @@ -1,10 +1,14 @@ # coding: utf- from __future__ import unicode_literals, print_function + +import signal + from future.builtins import str as text import os import requests import threadpool import time +import multiprocessing as mp try: from urllib.parse import urlparse @@ -13,13 +17,13 @@ except ImportError: from nhentai.logger import logger from nhentai.parser import request -from nhentai.utils import Singleton - +from nhentai.utils import Singleton, signal_handler requests.packages.urllib3.disable_warnings() +semaphore = mp.Semaphore() -class NhentaiImageNotExistException(Exception): +class NHentaiImageNotExistException(Exception): pass @@ -28,14 +32,14 @@ class Downloader(Singleton): def __init__(self, path='', thread=1, timeout=30, delay=0): 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.thread_pool = None self.timeout = timeout self.delay = delay - def _download(self, url, folder='', filename='', retried=0): + def download_(self, url, folder='', filename='', retried=0): if self.delay: time.sleep(self.delay) logger.info('Starting to download {0} ...'.format(url)) @@ -54,9 +58,9 @@ class Downloader(Singleton): try: response = request('get', url, stream=True, timeout=self.timeout) if response.status_code != 200: - raise NhentaiImageNotExistException + raise NHentaiImageNotExistException - except NhentaiImageNotExistException as e: + except NHentaiImageNotExistException as e: raise e except Exception as e: @@ -78,27 +82,37 @@ class Downloader(Singleton): except (requests.HTTPError, requests.Timeout) as e: 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) + return 0, self.download_(url=url, folder=folder, filename=filename, retried=retried+1) else: return 0, None - except NhentaiImageNotExistException as e: + except NHentaiImageNotExistException as e: os.remove(os.path.join(folder, base_filename.zfill(3) + extension)) return -1, url except Exception as e: + import traceback + traceback.print_stack() logger.critical(str(e)) return 0, None + except KeyboardInterrupt: + return -3, None + return 1, url - def _download_callback(self, request, result): + def _download_callback(self, result): result, data = result if result == 0: logger.warning('fatal errors occurred, ignored') # exit(1) elif result == -1: logger.warning('url {} return status code 404'.format(data)) + elif result == -2: + logger.warning('Ctrl-C pressed, exiting sub processes ...') + elif result == -3: + # workers wont be run, just pass + pass else: logger.log(15, '{0} downloaded successfully'.format(data)) @@ -119,10 +133,31 @@ class Downloader(Singleton): else: logger.warn('Path \'{0}\' already exist.'.format(folder)) - queue = [([url], {'folder': folder}) for url in queue] + queue = [(self, url, folder) for url in queue] - self.thread_pool = threadpool.ThreadPool(self.thread_count) - requests_ = threadpool.makeRequests(self._download, queue, self._download_callback) - [self.thread_pool.putRequest(req) for req in requests_] + pool = mp.Pool(self.thread_count, init_worker) + for item in queue: + pool.apply_async(download_wrapper, args=item, callback=self._download_callback) - self.thread_pool.wait() + pool.close() + print(1) + pool.join() + print(2) + + +def download_wrapper(obj, url, folder=''): + if semaphore.get_value(): + return Downloader.download_(obj, url=url, folder=folder) + else: + return -3, None + + +def init_worker(): + signal.signal(signal.SIGINT, subprocess_signal) + + +def subprocess_signal(signal, frame): + if semaphore.acquire(timeout=1): + logger.warning('Ctrl-C pressed, exiting sub processes ...') + + raise KeyboardInterrupt diff --git a/nhentai/utils.py b/nhentai/utils.py index 6a0d00b..815c690 100644 --- a/nhentai/utils.py +++ b/nhentai/utils.py @@ -207,3 +207,10 @@ an invalid filename. # Remove [] from filename filename = filename.replace('[]', '') return filename + + +def signal_handler(signal, frame): + logger.error('Ctrl-C signal received. Stopping...') + exit(1) + +