diff --git a/main.py b/main.py index 585c3e6..8578c75 100644 --- a/main.py +++ b/main.py @@ -7,7 +7,7 @@ from telethon import TelegramClient from telethon.tl import types from tools.down_file import down_group -from tools.tool import print_all_channel,Hook +from tools.tool import print_all_channel, Hook from tools.upload_file import upload_file config_path = './config.json' @@ -83,5 +83,5 @@ if __name__ == '__main__': client.loop.run_until_complete(down_group(client, channel_id, plus_func)) elif select == '3': channel_id = input('上传到:') - folder_path = input('文件夹路径:') + folder_path = input('文件(夹)路径:') client.loop.run_until_complete(upload_file(client, channel_id, folder_path)) diff --git a/tools/down_file.py b/tools/down_file.py index 9c6f16d..7eb6b95 100644 --- a/tools/down_file.py +++ b/tools/down_file.py @@ -54,8 +54,7 @@ async def download_file(channel_title, channel_id, message): download_path = file_path + '.downloading' print(f"开始下载:{file_name}") try: - with TqdmUpTo(unit='B', unit_scale=True, unit_divisor=1024, total=file_size, - bar_format=TqdmUpTo.bar_format, desc=file_name[:10]) as bar: + with TqdmUpTo(total=file_size, bar_format=TqdmUpTo.bar_format, desc=file_name[:10]) as bar: await message.download_media(download_path, progress_callback=bar.update_to) except CancelledError: print("取消下载") diff --git a/tools/tqdm.py b/tools/tqdm.py index 4682e03..1d96020 100644 --- a/tools/tqdm.py +++ b/tools/tqdm.py @@ -4,9 +4,15 @@ from tqdm import tqdm class TqdmUpTo(tqdm): total = None now_size = 0 - bar_format = '{l_bar}{bar}| {n_fmt}/{total_fmt} [已用时:{elapsed}预计剩余:{remaining}, {rate_fmt}{postfix}]' + def __init__(self, **kwargs): + super().__init__(**kwargs) + self.unit = 'B' + self.unit_scale = True + self.unit_divisor = 1024 + self.bar_format = TqdmUpTo.bar_format + def update_to(self, current, total): """更新进度条 :param current: 已传输 diff --git a/tools/upload_file.py b/tools/upload_file.py index 5a96c6a..9478e67 100644 --- a/tools/upload_file.py +++ b/tools/upload_file.py @@ -7,22 +7,33 @@ from telethon import TelegramClient from tools.tqdm import TqdmUpTo -async def upload_file(client: TelegramClient, chat_id, folder_path: str): +async def upload_file(client: TelegramClient, chat_id, path: str): peo = await client.get_entity(chat_id) + path_list = [] + if os.path.isfile(path): + path_list.append(path) + else: + for filename in os.listdir(path): + file_path = os.path.join(path, filename) + if os.path.isfile(file_path): + path_list.append(file_path) # 遍历文件夹下的所有文件 - for filename in os.listdir(folder_path): - file_path = os.path.join(folder_path, filename) - # 判断文件是否为视频文件 - if os.path.isfile(file_path): - try: - with TqdmUpTo(unit='B', unit_scale=True, unit_divisor=1024, total=os.path.getsize(file_path), - bar_format=TqdmUpTo.bar_format, desc=os.path.basename(file_path)) as bar: - # 上传文件到Telegram服务器 - result = await client.upload_file(file_path, progress_callback=bar.update_to) - # 发送文件到指定的群组或频道 - await client.send_file(peo, result, caption=os.path.basename(file_path)) - except CancelledError: - print("取消上传") - sys.exit() - except Exception as e: - print("上传出错", e.__class__.__name__) + for file_path in path_list: + # 文件预处理,解析信息 + if os.path.splitext(file_path)[1].lower() == '.mp4': + supports_streaming = True + else: + supports_streaming = False + file_caption = os.path.basename(file_path) + file_size = os.path.getsize(file_path) + # 发送文件到指定的群组或频道 + try: + with TqdmUpTo(total=file_size, desc=file_caption) as bar: + # 上传文件到Telegram服务器 + result = await client.upload_file(file_path, progress_callback=bar.update_to) + await client.send_file(peo, result, caption=file_caption, supports_streaming=supports_streaming) + except CancelledError: + print("取消上传") + sys.exit() + except Exception as e: + print("上传出错", e.__class__.__name__)