diff --git a/DownloadKit/_funcs.py b/DownloadKit/_funcs.py index a91e29daf8657c37a05bb11908553cceb799ce9a..8ec6be92d3ed18af6773636b5102af1c75d86d8a 100644 --- a/DownloadKit/_funcs.py +++ b/DownloadKit/_funcs.py @@ -13,6 +13,9 @@ from urllib.parse import unquote from requests import Session +FILE_EXISTS_MODE = {'rename': 'rename', 'overwrite': 'overwrite', 'skip': 'skip', 'add': 'add', 'r': 'rename', + 'o': 'overwrite', 's': 'skip', 'a': 'add'} + def copy_session(session): """复制输入Session对象,返回一个新的 @@ -67,15 +70,23 @@ class PathSetter(object): class FileExistsSetter(object): def __set__(self, file_exists, mode): - mode = mode.lower() - if mode not in ('skip', 'overwrite', 'rename', 'add'): - raise ValueError("file_exists参数只能传入'skip', 'overwrite', 'rename', 'add'") - file_exists._file_exists = mode + file_exists._file_exists = get_file_exists_mode(mode) def __get__(self, file_exists, objtype=None): return file_exists._file_exists +def get_file_exists_mode(mode): + """获取文件重名时处理策略名称 + :param mode: 输入 + :return: 标准字符串 + """ + mode = FILE_EXISTS_MODE.get(mode, mode) + if mode not in FILE_EXISTS_MODE: + raise ValueError(f'''mode参数只能是 '{"', '".join(FILE_EXISTS_MODE.keys())}' 之一,现在是:{mode}''') + return mode + + def get_usable_path(path): """检查文件或文件夹是否有重名,并返回可以使用的路径 :param path: 文件或文件夹路径 @@ -164,12 +175,13 @@ def set_charset(response): return response -def get_file_info(response, goal_path=None, rename=None, file_exists=None, lock=None): +def get_file_info(response, goal_path=None, rename=None, suffix=None, file_exists=None, lock=None): """获取文件信息,大小单位为byte 包括:size、path、skip :param response: Response对象 :param goal_path: 目标文件夹 :param rename: 重命名 + :param suffix: 重命名后缀名 :param file_exists: 存在重名文件时的处理方式 :param lock: 线程锁 :return: 文件名、文件大小、保存路径、是否跳过 @@ -189,13 +201,23 @@ def get_file_info(response, goal_path=None, rename=None, file_exists=None, lock= goal_Path.mkdir(parents=True, exist_ok=True) # ------------获取保存文件名------------ - # -------------------重命名,不改变扩展名------------------- + # -------------------重命名------------------- if rename: - tmp = file_name.rsplit('.', 1) - ext_name = tmp[-1] if len(tmp) > 1 else '' - tmp = rename.rsplit('.', 1) - ext_rename = tmp[-1] if len(tmp) > 1 else '' - full_name = rename if ext_rename == ext_name else f'{rename}.{ext_name}' + if suffix is not None: + full_name = f'{rename}.{suffix}' if suffix else rename + + else: + tmp = file_name.rsplit('.', 1) + ext_name = tmp[-1] if len(tmp) > 1 else '' + tmp = rename.rsplit('.', 1) + ext_rename = tmp[-1] if len(tmp) > 1 else '' + full_name = rename if ext_rename == ext_name else f'{rename}.{ext_name}' + + elif suffix is not None: + full_name = file_name.rsplit(".", 1)[0] + if suffix: + full_name = f'{full_name}.{suffix}' + else: full_name = file_name diff --git a/DownloadKit/_funcs.pyi b/DownloadKit/_funcs.pyi index 6f3ab2b6472750f69814143e82f516332561da40..2f57c15ce50f824d207d1d5c5ad1148bd4e82168 100644 --- a/DownloadKit/_funcs.pyi +++ b/DownloadKit/_funcs.pyi @@ -9,6 +9,8 @@ from typing import Union from requests import Session, Response +FILE_EXISTS_MODE: dict = ... + def copy_session(session: Session) -> Session: ... @@ -31,6 +33,9 @@ class FileExistsSetter(object): def __get__(self, file_exists, objtype=None): ... +def get_file_exists_mode(mode: str) -> str: ... + + def get_usable_path(path: Union[str, Path]) -> Path: ... @@ -43,11 +48,8 @@ def get_long(txt: str) -> int: ... def set_charset(response: Response) -> Response: ... -def get_file_info(response: Response, - goal_path: str = None, - rename: str = None, - file_exists: str = None, - lock: Lock = None) -> dict: ... +def get_file_info(response: Response, goal_path: str = None, rename: str = None, suffix: str = None, + file_exists: str = None, lock: Lock = None) -> dict: ... def set_session_cookies(session: Session, cookies: list) -> None: ... diff --git a/DownloadKit/downloadKit.py b/DownloadKit/downloadKit.py index f35aadc5df874980c72f9cc28e3a0407e532e0d9..630b952bb4ab0f9aa04855520d9c0e87673d14f3 100644 --- a/DownloadKit/downloadKit.py +++ b/DownloadKit/downloadKit.py @@ -14,7 +14,7 @@ from time import sleep, perf_counter from requests import Response from requests.structures import CaseInsensitiveDict -from ._funcs import FileExistsSetter, PathSetter, BlockSizeSetter, set_charset, get_file_info +from ._funcs import FileExistsSetter, PathSetter, BlockSizeSetter, set_charset, get_file_info, get_file_exists_mode from .mission import Task, Mission from .setter import Setter @@ -131,11 +131,12 @@ class DownloadKit(object): """用list方式返回所有任务对象""" return self._missions - def add(self, file_url, goal_path=None, rename=None, file_exists=None, split=None, **kwargs): + def add(self, file_url, goal_path=None, rename=None, suffix=None, file_exists=None, split=None, **kwargs): """添加一个下载任务并将其返回 :param file_url: 文件网址 :param goal_path: 保存路径 :param rename: 重命名的文件名 + :param suffix: 重命名的文件后缀名 :param file_exists: 遇到同名文件时的处理方式,可选 'skip', 'overwrite', 'rename', 'add',默认跟随实例属性 :param split: 是否允许多线程分块下载,为None则使用对象属性 :param kwargs: 连接参数 @@ -143,20 +144,21 @@ class DownloadKit(object): """ self._missions_num += 1 self._running_count += 1 - mission = Mission(self._missions_num, self, file_url, - str(goal_path or self.goal_path), - rename, file_exists or self.file_exists, - self.split if split is None else split, - kwargs) + if file_exists is not None: + file_exists = get_file_exists_mode(file_exists) + mission = Mission(self._missions_num, self, file_url, str(goal_path or self.goal_path), + rename, suffix, file_exists or self.file_exists, + self.split if split is None else split, kwargs) self._missions[self._missions_num] = mission self._run_or_wait(mission) return mission - def download(self, file_url, goal_path=None, rename=None, file_exists=None, show_msg=True, **kwargs): + def download(self, file_url, goal_path=None, rename=None, suffix=None, file_exists=None, show_msg=True, **kwargs): """以阻塞的方式下载一个文件并返回结果 :param file_url: 文件网址 :param goal_path: 保存路径 :param rename: 重命名的文件名 + :param suffix: 重命名的文件后缀名 :param file_exists: 遇到同名文件时的处理方式,可选 'skip', 'overwrite', 'rename', 'add',默认跟随实例属性 :param show_msg: 是否打印进度 :param kwargs: 连接参数 @@ -165,7 +167,7 @@ class DownloadKit(object): if show_msg: tmp = self._print_mode self._print_mode = None - r = self.add(file_url=file_url, goal_path=goal_path, rename=rename, file_exists=file_exists, + r = self.add(file_url=file_url, goal_path=goal_path, rename=rename, suffix=suffix, file_exists=file_exists, split=False, **kwargs).wait(show=show_msg) if show_msg: self._print_mode = tmp @@ -392,6 +394,7 @@ class DownloadKit(object): self._logger.add_data(('开始下载', mission.data.url)) rename = mission.data.rename + suffix = mission.data.suffix goal_path = mission.data.goal_path file_exists = mission.data.file_exists split = mission.data.split @@ -403,11 +406,13 @@ class DownloadKit(object): goal_Path.mkdir(parents=True, exist_ok=True) goal_path = str(goal_Path) - if file_exists == 'skip' and rename and (goal_Path / rename).exists(): - mission.file_name = rename - mission._set_path(goal_Path / rename) - mission._set_done('skipped', str(mission.path)) - return + if file_exists == 'skip' and rename: + tmp = goal_Path / rename + if tmp.exists() and tmp.is_file(): + mission.file_name = rename + mission._set_path(goal_Path / rename) + mission._set_done('skipped', str(mission.path)) + return r, inf = self._connect(file_url, mission.session, mission.method, **kwargs) @@ -419,7 +424,7 @@ class DownloadKit(object): return # -------------------获取文件信息------------------- - file_info = get_file_info(r, goal_path, rename, file_exists, self._lock) + file_info = get_file_info(r, goal_path, rename, suffix, file_exists, self._lock) file_size = file_info['size'] full_path = file_info['path'] mission._set_path(full_path) diff --git a/DownloadKit/downloadKit.pyi b/DownloadKit/downloadKit.pyi index 7cc1c7153a5c07d8b45aee1241172e6fc36d8b7d..de3c5413d7b654db10229e80954c6146b1c7e4a3 100644 --- a/DownloadKit/downloadKit.pyi +++ b/DownloadKit/downloadKit.pyi @@ -16,7 +16,7 @@ from ._funcs import FileExistsSetter, PathSetter, BlockSizeSetter from .mission import Task, Mission, BaseTask from .setter import Setter -FILE_EXISTS = Literal['add', 'skip', 'rename', 'overwrite'] +FILE_EXISTS = Literal['add', 'skip', 'rename', 'overwrite', 'a', 's', 'r', 'o'] class DownloadKit(object): @@ -102,6 +102,7 @@ class DownloadKit(object): file_url: str, goal_path: Optional[str, Path] = None, rename: str = None, + suffix: str = None, file_exists: FILE_EXISTS = None, split: bool = None, timeout: Optional[float] = None, @@ -123,6 +124,7 @@ class DownloadKit(object): file_url: str, goal_path: Optional[str, Path] = None, rename: str = None, + suffix: str = None, file_exists: FILE_EXISTS = None, show_msg: bool = True, timeout: Optional[float] = None, diff --git a/DownloadKit/mission.py b/DownloadKit/mission.py index 270557840c3bb4d8ff4242ba75110b33ef57f69d..92f32f42b8264c1982726a4cc78bb942b321f575 100644 --- a/DownloadKit/mission.py +++ b/DownloadKit/mission.py @@ -14,19 +14,21 @@ from ._funcs import copy_session, set_session_cookies class MissionData(object): - def __init__(self, url, goal_path, rename, file_exists, split, kwargs, offset=0): + def __init__(self, url, goal_path, rename, suffix, file_exists, split, kwargs, offset=0): """保存任务数据的对象 :param url: 下载文件url :param goal_path: 保存文件夹 :param rename: 文件重命名 + :param suffix: 文件重命名后缀名 :param file_exists: 存在重名文件时处理方式 :param split: 是否允许分块下载 :param kwargs: requests其它参数 :param offset: 文件存储偏移量 """ - self.url = quote(url, safe='/:&?=%;#@+![]') + self.url = quote(url, safe='-_.~!*\'();:@&=+$,/?#[]') self.goal_path = goal_path self.rename = rename + self.suffix = suffix self.file_exists = file_exists self.split = split self.kwargs = kwargs @@ -74,14 +76,14 @@ class BaseTask(object): class Mission(BaseTask): - def __init__(self, ID, download_kit, file_url, goal_path, rename, - file_exists, split, kwargs): + def __init__(self, ID, download_kit, file_url, goal_path, rename, suffix, file_exists, split, kwargs): """任务类 :param ID: 任务id :param download_kit: 所属DownloadKit对象 :param file_url: 文件网址 :param goal_path: 保存文件夹路径 :param rename: 重命名 + :param suffix: 重命名后缀名 :param file_exists: 存在同名文件处理方式 :param split: 是否分块下载 :param kwargs: 连接参数 @@ -100,7 +102,7 @@ class Mission(BaseTask): self.session = self._set_session() kwargs = self._handle_kwargs(file_url, kwargs) - self._data = MissionData(file_url, goal_path, rename, file_exists, split, kwargs) + self._data = MissionData(file_url, goal_path, rename, suffix, file_exists, split, kwargs) self.method = 'post' if (self._data.kwargs.get('data', None) is not None or self._data.kwargs.get('json', None) is not None) else 'get' diff --git a/DownloadKit/mission.pyi b/DownloadKit/mission.pyi index c91570fcd86250cf4ad73279418dfe0ce9b12e95..237fcf513f42f8cff9902e9960303e0816cd38b9 100644 --- a/DownloadKit/mission.pyi +++ b/DownloadKit/mission.pyi @@ -13,12 +13,13 @@ class MissionData(object): url: str = ... goal_path: Union[str, Path] = ... rename: Optional[str] = ... + suffix: Optional[str] = ... file_exists: str = ... split: bool = ... kwargs: dict = ... offset: int = ... - def __init__(self, url: str, goal_path: Union[str, Path], rename: Optional[str], + def __init__(self, url: str, goal_path: Union[str, Path], rename: Optional[str], suffix: Optional[str], file_exists: str, split: bool, kwargs: dict, offset: int = 0): ... @@ -64,7 +65,7 @@ class Mission(BaseTask): method: str = ... def __init__(self, ID: int, download_kit: DownloadKit, file_url: str, goal_path: Union[str, Path], rename: str, - file_exists: str, split: bool, kwargs: dict): ... + suffix: str, file_exists: str, split: bool, kwargs: dict): ... def __repr__(self) -> str: ... diff --git a/DownloadKit/setter.py b/DownloadKit/setter.py index 43fe163d22bc54317019ace71b8c76ced3422475..5f5a6314d7e7ae30a12b9bcf5235058c60d871b2 100644 --- a/DownloadKit/setter.py +++ b/DownloadKit/setter.py @@ -7,6 +7,8 @@ from DataRecorder import Recorder from requests import Session +from ._funcs import get_file_exists_mode + class Setter(object): def __init__(self, downloadKit): @@ -192,9 +194,7 @@ class FileExists(object): :param mode: 'skip', 'rename', 'overwrite', 'add' :return: None """ - if mode not in ('skip', 'rename', 'overwrite', 'add'): - raise ValueError("mode参数只能是'skip', 'rename', 'overwrite', 'add'") - self._setter._downloadKit.file_exists = mode + self._setter._downloadKit.file_exists = get_file_exists_mode(mode) def skip(self): """设为跳过""" diff --git a/DownloadKit/setter.pyi b/DownloadKit/setter.pyi index 359325be7101c7629126296d3f8cc7eff4b91b2e..42820c7e8ec9b783fecf4b54f57c468eca0afe73 100644 --- a/DownloadKit/setter.pyi +++ b/DownloadKit/setter.pyi @@ -6,13 +6,13 @@ from pathlib import Path from typing import Union, Literal -from DrissionPage.base import BasePage from DrissionPage import SessionOptions +from DrissionPage.base import BasePage from requests import Session from .downloadKit import DownloadKit -FILE_EXISTS = Literal['add', 'skip', 'rename', 'overwrite'] +FILE_EXISTS = Literal['add', 'skip', 'rename', 'overwrite', 'a', 's', 'r', 'o'] class Setter(object): diff --git a/setup.py b/setup.py index 639af0d0750fcecef822856fa84fad66cc9651d5..d97b2208a2103a007dd15ebfb21e26e6cb02ce77 100644 --- a/setup.py +++ b/setup.py @@ -6,7 +6,7 @@ with open("README.md", "r", encoding='utf-8') as fh: setup( name="DownloadKit", - version="1.0.3", + version="1.0.4", author="g1879", author_email="g1879@qq.com", description="一个简洁易用的多线程文件下载工具。",