Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
196 changes: 102 additions & 94 deletions DownloadKit/_funcs.py
Original file line number Diff line number Diff line change
Expand Up @@ -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对象,返回一个新的
Expand Down Expand Up @@ -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(';'):
Expand All @@ -164,45 +121,59 @@ 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

full_name = make_valid_name(full_name)

# -------------------生成路径-------------------
skip = False
overwrite = False
create = True
full_path = goal_Path / full_name

Expand All @@ -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':
Expand All @@ -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("'")

Expand All @@ -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):
Expand Down
Loading