This commit is contained in:
ricterz 2024-12-11 23:47:48 +08:00
parent 1f3528afad
commit b841747761

View File

@ -4,6 +4,7 @@ import os
import asyncio import asyncio
import httpx import httpx
import urllib3.exceptions import urllib3.exceptions
import math
from urllib.parse import urlparse from urllib.parse import urlparse
from nhentai import constant from nhentai import constant
@ -13,6 +14,7 @@ from nhentai.utils import Singleton, async_request
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning) urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
class NHentaiImageNotExistException(Exception): class NHentaiImageNotExistException(Exception):
pass pass
@ -32,13 +34,14 @@ def download_callback(result):
logger.log(16, f'{data} downloaded successfully') logger.log(16, f'{data} downloaded successfully')
class Downloader(Singleton): class Downloader(Singleton):
def __init__(self, path='', threads=5, timeout=30, delay=0): def __init__(self, path='', threads=5, timeout=30, delay=0):
self.threads = threads self.threads = threads
self.path = str(path) self.path = str(path)
self.timeout = timeout self.timeout = timeout
self.delay = delay self.delay = delay
self.folder = None
self.semaphore = None
async def fiber(self, tasks): async def fiber(self, tasks):
self.semaphore = asyncio.Semaphore(self.threads) self.semaphore = asyncio.Semaphore(self.threads)
@ -49,18 +52,20 @@ class Downloader(Singleton):
except Exception as e: except Exception as e:
logger.error(f'An error occurred: {e}') logger.error(f'An error occurred: {e}')
async def _semaphore_download(self, *args, **kwargs): async def _semaphore_download(self, *args, **kwargs):
async with self.semaphore: async with self.semaphore:
return await self.download(*args, **kwargs) return await self.download(*args, **kwargs)
async def download(self, url, folder='', filename='', retried=0, proxy=None): async def download(self, url, folder='', filename='', retried=0, proxy=None, length=0):
logger.info(f'Starting to download {url} ...') logger.info(f'Starting to download {url} ...')
if self.delay: if self.delay:
await asyncio.sleep(self.delay) await asyncio.sleep(self.delay)
filename = filename if filename else os.path.basename(urlparse(url).path) filename = filename if filename else os.path.basename(urlparse(url).path)
base_filename, extension = os.path.splitext(filename)
digits = int(math.log10(length)) + 1
filename = base_filename.zfill(digits) + extension
save_file_path = os.path.join(self.folder, filename) save_file_path = os.path.join(self.folder, filename)
@ -129,7 +134,6 @@ class Downloader(Singleton):
f.write(chunk) f.write(chunk)
return True return True
def start_download(self, queue, folder='') -> bool: def start_download(self, queue, folder='') -> bool:
if not isinstance(folder, (str,)): if not isinstance(folder, (str,)):
folder = str(folder) folder = str(folder)
@ -149,9 +153,8 @@ class Downloader(Singleton):
# Assuming we want to continue with rest of process. # Assuming we want to continue with rest of process.
return True return True
coroutines = [ coroutines = [
self._semaphore_download(url, filename=os.path.basename(urlparse(url).path)) self._semaphore_download(url, filename=os.path.basename(urlparse(url).path), length=len(queue))
for url in queue for url in queue
] ]