diff --git a/DownloadKit/_funcs.py b/DownloadKit/_funcs.py index a91e29d..b7f7856 100644 --- a/DownloadKit/_funcs.py +++ b/DownloadKit/_funcs.py @@ -7,12 +7,16 @@ from os import path as os_PATH from pathlib import Path from random import randint -from re import search, sub +from re import search, sub, IGNORECASE from time import time from urllib.parse import unquote +from DataRecorder.tools import get_usable_path, make_valid_name 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对象,返回一个新的 @@ -56,91 +60,44 @@ def __get__(self, block_size, objtype=None) -> int: class PathSetter(object): - def __set__(self, goal_path, val): + def __set__(self, save_path, val): if val is not None and not isinstance(val, (str, Path)): - raise TypeError('goal_path只能是str或Path类型。') - goal_path._goal_path = str(val) if isinstance(val, Path) else val + raise TypeError('路径只能是str或Path类型。') + save_path._save_path = str(val) if isinstance(val, Path) else val - def __get__(self, goal_path, objtype=None): - return goal_path._goal_path + def __get__(self, save_path, objtype=None): + return save_path._save_path 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_usable_path(path): - """检查文件或文件夹是否有重名,并返回可以使用的路径 - :param path: 文件或文件夹路径 - :return: 可用的路径,Path对象 - """ - path = Path(path) - parent = path.parent - path = parent / make_valid_name(path.name) - name = path.stem if path.is_file() else path.name - ext = path.suffix if path.is_file() else '' - - first_time = True - - while path.exists(): - r = search(r'(.*)_(\d+)$', name) - - if not r or (r and first_time): - src_name, num = name, '1' - else: - src_name, num = r.group(1), int(r.group(2)) + 1 - - name = f'{src_name}_{num}' - path = parent / f'{name}{ext}' - first_time = None - - return path - - -def make_valid_name(full_name): - """获取有效的文件名 - :param full_name: 文件名 - :return: 可用的文件名 +def get_file_exists_mode(mode): + """获取文件重名时处理策略名称 + :param mode: 输入 + :return: 标准字符串 """ - # ----------------去除前后空格---------------- - full_name = full_name.strip() - - # ----------------使总长度不大于255个字符(一个汉字是2个字符)---------------- - r = search(r'(.*)(\.[^.]+$)', full_name) # 拆分文件名和后缀名 - if r: - name, ext = r.group(1), r.group(2) - ext_long = len(ext) - else: - name, ext = full_name, '' - ext_long = 0 - - while get_long(name) > 255 - ext_long: - name = name[:-1] + 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 - full_name = f'{name}{ext}' - # ----------------去除不允许存在的字符---------------- - return sub(r'[<>/\\|:*?\n]', '', full_name) - - -def get_long(txt): - """返回字符串中字符个数(一个汉字是2个字符) - :param txt: 字符串 - :return: 字符个数 +def set_charset(response, encoding): + """设置Response对象的编码 + :param response: Response对象 + :param encoding: 指定的编码格式 + :return: 设置编码后的Response对象 """ - txt_len = len(txt) - return int((len(txt.encode('utf-8')) - txt_len) / 2 + txt_len) + if encoding: + response.encoding = encoding + return response - -def set_charset(response): - """设置Response对象的编码""" # 在headers中获取编码 content_type = response.headers.get('content-type', '').lower() if not content_type.endswith(';'): @@ -164,38 +121,51 @@ 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, save_path=None, rename=None, suffix=None, file_exists=None, encoding=None, lock=None): """获取文件信息,大小单位为byte 包括:size、path、skip :param response: Response对象 - :param goal_path: 目标文件夹 + :param save_path: 目标文件夹 :param rename: 重命名 + :param suffix: 重命名后缀名 :param file_exists: 存在重名文件时的处理方式 + :param encoding: 编码格式 :param lock: 线程锁 - :return: 文件名、文件大小、保存路径、是否跳过 + :return: 文件大小、完整路径、是否跳过、是否覆盖 """ # ------------获取文件大小------------ file_size = response.headers.get('Content-Length', None) file_size = None if file_size is None else int(file_size) # ------------获取网络文件名------------ - file_name = _get_file_name(response) + file_name = _get_file_name(response, encoding) # ------------获取保存路径------------ - goal_Path = Path(goal_path) + goal_Path = Path(save_path) # 按windows规则去除路径中的非法字符 - goal_path = goal_Path.anchor + sub(r'[*:|<>?"]', '', goal_path.lstrip(goal_Path.anchor)).strip() - goal_Path = Path(goal_path).absolute() + g = save_path[len(goal_Path.anchor):] if save_path.lower().startswith(goal_Path.anchor.lower()) else save_path + save_path = goal_Path.anchor + sub(r'[*:|<>?"]', '', g).strip() + goal_Path = Path(save_path).absolute() 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 = f'.{tmp[-1]}' if len(tmp) > 1 else '' + tmp = rename.rsplit('.', 1) + ext_rename = f'.{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 @@ -203,6 +173,7 @@ def get_file_info(response, goal_path=None, rename=None, file_exists=None, lock= # -------------------生成路径------------------- skip = False + overwrite = False create = True full_path = goal_Path / full_name @@ -216,6 +187,7 @@ def get_file_info(response, goal_path=None, rename=None, file_exists=None, lock= create = False elif file_exists == 'overwrite': + overwrite = True full_path.unlink() elif file_exists == 'add': @@ -227,35 +199,42 @@ def get_file_info(response, goal_path=None, rename=None, file_exists=None, lock= return {'size': file_size, 'path': full_path, - 'skip': skip} + 'skip': skip, + 'overwrite': overwrite} -def _get_file_name(response) -> str: +def _get_file_name(response, encoding) -> str: """从headers或url中获取文件名,如果获取不到,生成一个随机文件名 :param response: 返回的response + :param encoding: 在headers获取时指定编码格式 :return: 下载文件的文件名 """ file_name = '' charset = '' - content_disposition = response.headers.get('content-disposition', '').replace(' ', '') + content_disposition = response.headers.get('content-disposition', '') # 使用header里的文件名 if content_disposition: - txt = search(r'filename\*="?([^";]+)', content_disposition) + # 先尝试匹配 filename*= 格式(RFC 5987) + txt = search(r'filename\*=([^;]+)', content_disposition, IGNORECASE) if txt: # 文件名自带编码方式 - txt = txt.group(1).split("''", 1) + # 移除可能的引号和空格 + filename_value = txt.group(1).strip(' "') + txt = filename_value.split("''", 1) if len(txt) == 2: charset, file_name = txt else: file_name = txt[0] else: # 文件名没带编码方式 - txt = search(r'filename="?([^";]+)', content_disposition) + # 匹配 filename="..." 或 filename=... 格式 + # 支持引号内的内容和无引号的内容 + txt = search(r'filename=([^;]+)', content_disposition, IGNORECASE) if txt: - file_name = txt.group(1) - - # 获取编码(如有) - charset = response.encoding + file_name = txt.group(1).strip(' "') + # 获取编码:如果用户指定了 encoding,使用用户指定的;否则使用 response.encoding + # 注意:这里 charset 用于 URL 解码,实际文件名解码会在后面用 encoding 参数处理 + charset = response.encoding or 'utf-8' file_name = file_name.strip("'") @@ -269,7 +248,36 @@ def _get_file_name(response) -> str: # 去除非法字符 charset = charset or 'utf-8' - return unquote(file_name, charset) + + # 先尝试 URL 解码 + decoded_name = unquote(file_name, charset) + + # 如果指定了 encoding,尝试按指定编码重新解码 + if encoding: + try: + # HTTP 响应头默认使用 ISO-8859-1 (latin-1) 编码 + # 如果文件名没有被 URL 编码(unquote 没有改变), + # 可能是响应头中的文件名已经是编码后的字节被当作 ISO-8859-1 字符串了 + # 先按 latin-1 编码成字节,然后按指定编码解码 + if decoded_name == file_name: + # 文件名没有被 URL 编码,尝试按 latin-1 编码再按指定编码解码 + file_name = decoded_name.encode('latin-1').decode(encoding) + else: + # 文件名已经 URL 解码,可能需要按指定编码重新处理 + # 如果 URL 解码后的字符串看起来像乱码,尝试重新解码 + try: + # 尝试将 URL 解码后的字符串按 latin-1 编码再按指定编码解码 + file_name = decoded_name.encode('latin-1').decode(encoding) + except (UnicodeEncodeError, UnicodeDecodeError): + # 如果失败,使用 URL 解码后的结果 + file_name = decoded_name + except (UnicodeEncodeError, UnicodeDecodeError): + # 如果解码失败,使用 URL 解码后的结果 + file_name = decoded_name + else: + file_name = decoded_name + + return file_name def set_session_cookies(session, cookies): diff --git a/DownloadKit/downloadKit.py b/DownloadKit/downloadKit.py index dfffbab..37c401f 100644 --- a/DownloadKit/downloadKit.py +++ b/DownloadKit/downloadKit.py @@ -5,6 +5,7 @@ @File : downloadKit.py """ from copy import copy +from os.path import sep from pathlib import Path from queue import Queue from re import sub @@ -14,22 +15,23 @@ 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 class DownloadKit(object): file_exists = FileExistsSetter() - goal_path = PathSetter() + save_path = PathSetter() block_size = BlockSizeSetter() - def __init__(self, goal_path=None, roads=10, session=None, file_exists='rename', driver=None): + def __init__(self, save_path=None, roads=10, driver=None, + file_exists='rename', session=None, goal_path=None): # session和goal_path即将废弃 """ - :param goal_path: 文件保存路径 + :param save_path: 文件保存路径 :param roads: 可同时运行的线程数 - :param file_exists: 有同名文件名时的处理方式,可选 'skip', 'overwrite', 'rename', 'add' :param driver: 使用的Session对象,或配置对象、页面对象等 + :param file_exists: 有同名文件名时的处理方式,可选 'skip', 'overwrite', 'rename', 'add', 's', 'o', 'r', 'a' """ self._roads = roads self._missions = {} @@ -43,31 +45,34 @@ def __init__(self, goal_path=None, roads=10, session=None, file_exists='rename', self._retry = None self._interval = None self._timeout = None - self._copy_cookies = False + self._encoding = None self._setter = None self._print_mode = None self._log_mode = None self._logger = None - self.goal_path = goal_path or '.' + self.save_path = goal_path or save_path or '.' # goal_path即将废弃 self.file_exists = file_exists self.split = True self.block_size = '50M' self.set.driver(session or driver) - def __call__(self, file_url, goal_path=None, rename=None, file_exists=None, show_msg=True, **kwargs): + def __call__(self, file_url, save_path=None, rename=None, suffix=None, + file_exists=None, show_msg=True, goal_path=None, **kwargs): # goal_path即将废弃 """以阻塞的方式下载一个文件并返回结果 :param file_url: 文件网址 - :param goal_path: 保存路径 + :param save_path: 保存路径 :param rename: 重命名的文件名 - :param file_exists: 遇到同名文件时的处理方式,可选 'skip', 'overwrite', 'rename', 'add',默认跟随实例属性 + :param suffix: 指定后缀名 + :param file_exists: 遇到同名文件时的处理方式,可选 'skip', 'overwrite', 'rename', 'add', 's', 'o', 'r', 'a',默认跟随实例属性 :param show_msg: 是否打印进度 :param kwargs: 连接参数 :return: 任务结果和信息组成的tuple """ - return self.download(file_url=file_url, goal_path=goal_path, rename=rename, file_exists=file_exists, - show_msg=show_msg, **kwargs) + save_path = goal_path or save_path # goal_path即将废弃 + return self.download(file_url=file_url, save_path=save_path, rename=rename, suffix=suffix, + file_exists=file_exists, show_msg=show_msg, **kwargs) @property def set(self): @@ -131,33 +136,43 @@ def missions(self): """用list方式返回所有任务对象""" return self._missions - def add(self, file_url, goal_path=None, rename=None, file_exists=None, split=None, **kwargs): + @property + def encoding(self): + """返回指定的编码格式""" + return self._encoding + + def add(self, file_url, save_path=None, rename=None, suffix=None, + file_exists=None, split=None, goal_path=None, **kwargs): # goal_path即将废弃 """添加一个下载任务并将其返回 :param file_url: 文件网址 - :param goal_path: 保存路径 + :param save_path: 保存路径 :param rename: 重命名的文件名 - :param file_exists: 遇到同名文件时的处理方式,可选 'skip', 'overwrite', 'rename', 'add',默认跟随实例属性 + :param suffix: 重命名的文件后缀名 + :param file_exists: 遇到同名文件时的处理方式,可选 'skip', 'overwrite', 'rename', 'add', 's', 'o', 'r', 'a',默认跟随实例属性 :param split: 是否允许多线程分块下载,为None则使用对象属性 :param kwargs: 连接参数 :return: 任务对象 """ 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, + if file_exists is not None: + file_exists = get_file_exists_mode(file_exists) + save_path = goal_path or save_path # goal_path即将废弃 + mission = Mission(self._missions_num, self, file_url, str(save_path or self.save_path), rename, suffix, + file_exists or self.file_exists, self.split if split is None else split, self._encoding, 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, save_path=None, rename=None, suffix=None, + file_exists=None, show_msg=True, goal_path=None, **kwargs): # goal_path即将废弃 """以阻塞的方式下载一个文件并返回结果 :param file_url: 文件网址 - :param goal_path: 保存路径 + :param save_path: 保存路径 :param rename: 重命名的文件名 - :param file_exists: 遇到同名文件时的处理方式,可选 'skip', 'overwrite', 'rename', 'add',默认跟随实例属性 + :param suffix: 重命名的文件后缀名 + :param file_exists: 遇到同名文件时的处理方式,可选 'skip', 'overwrite', 'rename', 'add', 's', 'o', 'r', 'a',默认跟随实例属性 :param show_msg: 是否打印进度 :param kwargs: 连接参数 :return: 任务结果和信息组成的tuple @@ -165,7 +180,8 @@ def download(self, file_url, goal_path=None, rename=None, file_exists=None, show 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, + save_path = goal_path or save_path # goal_path即将废弃 + r = self.add(file_url=file_url, save_path=save_path, rename=rename, suffix=suffix, file_exists=file_exists, split=False, **kwargs).wait(show=show_msg) if show_msg: self._print_mode = tmp @@ -234,7 +250,7 @@ def wait(self, mission=None, show=False, timeout=None): self.show(False) else: end_time = perf_counter() + timeout - while self.is_running or (perf_counter() < end_time or timeout == 0): + while self.is_running and (perf_counter() < end_time or timeout == 0): sleep(0.1) def cancel(self): @@ -288,18 +304,21 @@ def _show(self, wait, keep=False): print() - def _connect(self, url, session, method, **kwargs): + def _connect(self, url, session, _headers, method, encoding, **kwargs): """生成response对象 :param url: 目标url :param session: 用于连接的Session对象 + :param _headers: 内置的headers参数 :param method: 请求方式 + :param encoding: 编码格式 :param kwargs: 连接参数 :return: tuple,第一位为Response或None,第二位为出错信息或'Success' """ + kwargs = CaseInsensitiveDict(kwargs) if 'headers' in kwargs: # 不知道为什么添加这个才能正常使用多线程 - kwargs['headers'] = CaseInsensitiveDict(kwargs['headers']) + kwargs['headers'] = CaseInsensitiveDict({**_headers, **kwargs['headers']}) else: - kwargs['headers'] = CaseInsensitiveDict() + kwargs['headers'] = _headers r = err = None for i in range(self.retry + 1): @@ -310,7 +329,7 @@ def _connect(self, url, session, method, **kwargs): r = session.post(url, **kwargs) if r: - return set_charset(r), 'Success' + return set_charset(r, encoding), 'Success' except Exception as e: err = e @@ -349,11 +368,13 @@ def _when_mission_done(self, mission): if self._log_mode == 'all' or (self._log_mode == 'failed' and mission.result is False): data = ('下载结果', mission.data.url, - mission.data.goal_path, + mission.data.save_path, mission.data.rename, mission.data.kwargs) self._logger.add_data(data) + mission.session.close() + def _download(self, mission_or_task, thread_id): """此方法是执行下载的线程方法,用于根据任务下载文件 :param mission_or_task: 下载任务对象 @@ -372,7 +393,8 @@ def _download(self, mission_or_task, thread_id): kwargs = copy(mission_or_task.data.kwargs) task = mission_or_task kwargs['headers']['Range'] = f"bytes={task.range[0]}-{task.range[1]}" - r, inf = self._connect(file_url, task.mission.session, task.mission.method, **kwargs) + r, inf = self._connect(file_url, task.mission.session, task.mission.headers, + task.mission.method, task.mission.encoding, **kwargs) if r: _do_download(r, task, False) @@ -392,24 +414,27 @@ def _download(self, mission_or_task, thread_id): self._logger.add_data(('开始下载', mission.data.url)) rename = mission.data.rename - goal_path = mission.data.goal_path + suffix = mission.data.suffix + goal_Path = Path(mission.data.save_path) file_exists = mission.data.file_exists split = mission.data.split - goal_Path = Path(goal_path) - # 按windows规则去除路径中的非法字符 - goal_path = goal_Path.anchor + sub(r'[*:|<>?"]', '', goal_path.lstrip(goal_Path.anchor)).strip() - goal_Path = Path(goal_path).absolute() + parts = goal_Path.parts + if len(parts) > 1: # 按windows规则去除路径中的非法字符 + goal_Path = Path(goal_Path.anchor + sub(r'[*:|<>?"]', '', sep.join(parts[1:])).strip()) + goal_Path = goal_Path.absolute() goal_Path.mkdir(parents=True, exist_ok=True) - goal_path = str(goal_Path) + save_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) + r, inf = self._connect(file_url, mission.session, mission.headers, mission.method, mission.encoding, **kwargs) if mission.is_done: return @@ -419,12 +444,13 @@ def _download(self, mission_or_task, thread_id): return # -------------------获取文件信息------------------- - file_info = get_file_info(r, goal_path, rename, file_exists, self._lock) + file_info = get_file_info(r, save_path, rename, suffix, file_exists, mission.encoding, self._lock) file_size = file_info['size'] full_path = file_info['path'] mission._set_path(full_path) mission.file_name = full_path.name mission.size = file_size + mission._overwrote = file_info['overwrite'] if file_info['skip']: mission._set_done('skipped', str(mission.path))