From 3b8b95ad01b77a717f29974d7cc247ffd23c58b8 Mon Sep 17 00:00:00 2001 From: tangyoha Date: Sun, 14 Jul 2024 20:16:36 +0800 Subject: [PATCH] fix: update pyrogram --- README.md | 2 ++ README_CN.md | 3 ++- media_downloader.py | 30 ++++++++++++++++++++++++++++++ module/app.py | 15 +++++++++++---- module/pyrogram_extension.py | 2 +- module/send_media_group_v2.py | 2 -- requirements.txt | 4 ++-- tests/test_common.py | 4 ++-- tests/test_media_downloader.py | 30 ++++++++++++++++++++++++++++++ utils/__init__.py | 4 ++-- 10 files changed, 82 insertions(+), 14 deletions(-) diff --git a/README.md b/README.md index a41f8ca0..0d1d27ae 100644 --- a/README.md +++ b/README.md @@ -196,6 +196,7 @@ web_login_secret: 123 allowed_user_ids: - 'me' date_format: '%Y_%m' +enable_download_txt: false ``` - **api_hash** - The api_hash you got from telegram apps @@ -236,6 +237,7 @@ date_format: '%Y_%m' - **allowed_user_ids** - Who is allowed to use the robot? The default login account can be used. Please add single quotes to the name with @. - **date_format** Support custom configuration of media_datetime format in file_path_prefix.see [python-datetime](https://docs.python.org/3/library/datetime.html) - **drop_no_audio_video*** Since the telegram server does not support uploading a set of media, an error will be reported when uploading a video without audio, so use ffmpeg to process it. If not, set this option to true. The default is false. +- **enable_download_txt** Enable download txt file, default `false` ## Execution diff --git a/README_CN.md b/README_CN.md index 96595f61..c9898028 100644 --- a/README_CN.md +++ b/README_CN.md @@ -190,6 +190,7 @@ web_login_secret: 123 allowed_user_ids: - 'me' date_format: '%Y_%m' +enable_download_txt: false ``` - **api_hash** - 你从电报应用程序获得的 api_hash @@ -232,7 +233,7 @@ date_format: '%Y_%m' - **allowed_user_ids** - 允许哪些人使用机器人，默认登录账号可以使用，带@的名称请加单引号 - **date_format** - 支持自定义配置file_path_prefix中media_datetime的格式，具体格式查看 [python-datetime](https://docs.python.org/zh-cn/3/library/time.html) - **drop_no_audio_video** 由于telegram服务器不支持上传一组媒体的时候，上传没有音频的视频会报错，所以使用ffmpeg处理，如果不处理将该选项置为true，默认为false - +- **enable_download_txt** 启用下载txt文件，默认`false` ## 执行 diff --git a/media_downloader.py b/media_downloader.py index acff6b3f..3775272b 100644 --- a/media_downloader.py +++ b/media_downloader.py @@ -265,6 +265,33 @@ async def add_download_task( return True +async def save_msg_to_file( + app, chat_id: Union[int, str], message: pyrogram.types.Message +): + """Write message text into file""" + dirname = validate_title( + message.chat.title if message.chat and message.chat.title else str(chat_id) + ) + datetime_dir_name = message.date.strftime(app.date_format) if message.date else "0" + + file_save_path = app.get_file_save_path("msg", dirname, datetime_dir_name) + file_name = os.path.join( + app.temp_save_path, + file_save_path, + f"{app.get_file_name(message.id, None, None)}.txt", + ) + + os.makedirs(os.path.dirname(file_name), exist_ok=True) + + if _is_exist(file_name): + return DownloadStatus.SkipDownload, None + + with open(file_name, "w", encoding="utf-8") as f: + f.write(message.text or "") + + return DownloadStatus.SuccessDownload, file_name + + async def download_task( client: pyrogram.Client, message: pyrogram.types.Message, node: TaskNode ): @@ -274,6 +301,9 @@ async def download_task( client, message, app.media_types, app.file_formats, node ) + if app.enable_download_txt and message.text and not message.media: + download_status, file_name = await save_msg_to_file(app, node.chat_id, message) + if not node.bot: app.set_download_id(node, message.id, download_status) diff --git a/module/app.py b/module/app.py index a67d14f8..b576a4fb 100644 --- a/module/app.py +++ b/module/app.py @@ -391,6 +391,7 @@ def __init__( ) self.date_format: str = "%Y_%m" self.drop_no_audio_video: bool = False + self.enable_download_txt: bool = False self.forward_limit_call = LimitCall(max_limit_call_times=33) self.loop = asyncio.new_event_loop() @@ -473,14 +474,16 @@ def assign_config(self, _config: dict) -> bool: # TODO: add check if expression exist syntax error - self.max_concurrent_transmissions = _config.get( - "max_concurrent_transmissions", self.max_concurrent_transmissions - ) - self.max_download_task = _config.get( "max_download_task", self.max_download_task ) + self.max_concurrent_transmissions = self.max_download_task * 5 + + self.max_concurrent_transmissions = _config.get( + "max_concurrent_transmissions", self.max_concurrent_transmissions + ) + language = _config.get("language", "EN") try: @@ -520,6 +523,10 @@ def assign_config(self, _config: dict) -> bool: _config, "drop_no_audio_video", self.drop_no_audio_video, bool ) + self.enable_download_txt = get_config( + _config, "enable_download_txt", self.enable_download_txt, bool + ) + try: date = datetime(2023, 10, 31) date.strftime(self.date_format) diff --git a/module/pyrogram_extension.py b/module/pyrogram_extension.py index 68f56441..8a967dd6 100644 --- a/module/pyrogram_extension.py +++ b/module/pyrogram_extension.py @@ -254,7 +254,7 @@ async def upload_telegram_chat( """Upload telegram chat""" # upload telegram if node.upload_telegram_chat_id: - if download_status is DownloadStatus.SkipDownload: + if download_status is DownloadStatus.SkipDownload and message.media: if message.media_group_id: await proc_cache_forward(client, node, message, True) return diff --git a/module/send_media_group_v2.py b/module/send_media_group_v2.py index 05325e9f..806558b5 100644 --- a/module/send_media_group_v2.py +++ b/module/send_media_group_v2.py @@ -379,7 +379,6 @@ async def send_media_group_v2( chat_id: Union[int, str], multi_media: List[raw.types.InputSingleMedia], disable_notification: bool = None, - reply_to_message_id: int = None, schedule_date: datetime = None, protect_content: bool = None, ): @@ -391,7 +390,6 @@ async def send_media_group_v2( peer=await client.resolve_peer(chat_id), multi_media=multi_media, silent=disable_notification or None, - reply_to_msg_id=reply_to_message_id, schedule_date=utils.datetime_to_timestamp(schedule_date), noforwards=protect_content, ), diff --git a/requirements.txt b/requirements.txt index 95dbe710..2b52e908 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,6 +1,6 @@ #https://github.com/tangyoha/pyrogram/archive/refs/tags/v2.0.69.zip -#pyrogram==2.0.106 -https://github.com/tangyoha/pyrogram/archive/refs/heads/master_v2.0.69_1.zip +https://github.com/tangyoha/pyrogram/archive/refs/heads/patch.zip +#https://github.com/tangyoha/pyrogram/archive/refs/heads/master_v2.0.69_1.zip PyYAML==5.3.1 rich==12.5.1 TgCrypto==1.2.5 diff --git a/tests/test_common.py b/tests/test_common.py index e0fea422..9bbea4d3 100644 --- a/tests/test_common.py +++ b/tests/test_common.py @@ -30,7 +30,7 @@ def __init__(self, **kwargs): self.video_note = kwargs.get("video_note", None) self.media_group_id = kwargs.get("media_group_id", None) self.caption = kwargs.get("caption", None) - self.text = None + self.text = kwargs.get("text", None) self.empty = kwargs.get("empty", False) self.from_user = kwargs.get("from_user", None) self.reply_to_message_id = kwargs.get("reply_to_message_id", None) @@ -40,7 +40,7 @@ def __init__(self, **kwargs): kwargs.get("chat_id", None), kwargs.get("chat_title", None) ) else: - self.chat = None + self.chat = kwargs.get("chat", None) self.date: datetime = None if kwargs.get("date") != None: self.date = kwargs["date"] diff --git a/tests/test_media_downloader.py b/tests/test_media_downloader.py index e25cfb8e..bbdb6d71 100644 --- a/tests/test_media_downloader.py +++ b/tests/test_media_downloader.py @@ -21,6 +21,7 @@ download_media, download_task, main, + save_msg_to_file, worker, ) from module.app import Application, DownloadStatus, TaskNode @@ -886,6 +887,35 @@ def test_is_exist(self): result2 = _is_exist(this_dir) self.assertEqual(result2, False) + @mock.patch("media_downloader.os.makedirs") + @mock.patch("builtins.open", new_callable=mock.mock_open) + def test_save_msg_to_file(self, mock_open, mock_makedirs): + rest_app(MOCK_CONF) + app.enable_download_txt = True + app.temp_save_path = "/tmp" + app.date_format = "%Y_%m" + + message = MockMessage( + id=123, + dis_chat=True, + chat=Chat(chat_id=456, chat_title="Test Chat"), + date=datetime(2023, 5, 15, 10, 30, 0), + text="This is a test message", + ) + + expected_file_path = platform_generic_path( + "/root/project/Test Chat/2023_05/123.txt" + ) + + result = self.loop.run_until_complete(save_msg_to_file(app, 456, message)) + + self.assertEqual(result, (DownloadStatus.SuccessDownload, expected_file_path)) + mock_makedirs.assert_called_once_with( + os.path.dirname(expected_file_path), exist_ok=True + ) + mock_open.assert_called_once_with(expected_file_path, "w", encoding="utf-8") + mock_open().write.assert_called_once_with("This is a test message") + @mock.patch("media_downloader.RETRY_TIME_OUT", new=0) @mock.patch("media_downloader.os.path.getsize", new=os_get_file_size) @mock.patch("media_downloader.os.remove", new=os_remove) diff --git a/utils/__init__.py b/utils/__init__.py index 474a5de2..e6ab4b2e 100644 --- a/utils/__init__.py +++ b/utils/__init__.py @@ -1,5 +1,5 @@ """Init namespace""" -__version__ = "2.2.3" +__version__ = "2.2.4" __license__ = "MIT License" -__copyright__ = "Copyright (C) 2023 tangyoha " +__copyright__ = "Copyright (C) 2024 tangyoha "