fix: topic message can not download (#84)

* fix: topic message can not download * feat: support more filter
tangyoha · Jun 18, 2023 · 24a69ec · 24a69ec
1 parent 1e8a815
commit 24a69ec
Show file tree

Hide file tree

Showing 11 changed files with 167 additions and 106 deletions.
diff --git a/README.md b/README.md
@@ -282,4 +282,3 @@ Help us keep Telegram Media Downloader open and inclusive. Please read and follo
 <img alt="Code style: black" style="width:30%" src="./screenshot/alipay.JPG">
 <img alt="Code style: black" style="width:30%" src="./screenshot/wechat.JPG">
 </p>
-
diff --git a/media_downloader.py b/media_downloader.py
@@ -22,6 +22,7 @@
     record_download_status,
     report_bot_download_status,
     set_max_concurrent_transmissions,
+    set_meta_data,
     upload_telegram_chat,
 )
 from module.web import get_flask_app
@@ -520,9 +521,7 @@ async def download_chat_task(
             app.set_caption_name(node.chat_id, message.media_group_id, caption)
         else:
             caption = app.get_caption_name(node.chat_id, message.media_group_id)
-        meta_data.get_meta_data(message)
-        if caption:
-            meta_data.message_caption = caption
+        set_meta_data(meta_data, message, caption)
 
         if not app.need_skip_message(chat_download_config, message.id, meta_data):
             await add_download_task(message, node)

diff --git a/module/bot.py b/module/bot.py
@@ -24,6 +24,7 @@
     get_message_with_retry,
     report_bot_forward_status,
     report_bot_status,
+    set_meta_data,
 )
 from utils.format import extract_info_from_link, replace_date_time, validate_title
 from utils.meta_data import MetaData
@@ -738,7 +739,7 @@ async def forward_normal_content(
             _bot.app.set_caption_name(node.chat_id, message.media_group_id, caption)
         else:
             caption = _bot.app.get_caption_name(node.chat_id, message.media_group_id)
-        meta_data.get_meta_data(message)
+        set_meta_data(meta_data, message, caption)
         if not _bot.filter.exec(node.download_filter):
             forward_ret = ForwardStatus.SkipForward
             await report_bot_forward_status(client, node, forward_ret)

diff --git a/module/pyrogram_extension.py b/module/pyrogram_extension.py
@@ -27,6 +27,7 @@
 from module.download_stat import get_download_result
 from module.language import Language, _t
 from utils.format import create_progress_bar, format_byte, truncate_filename
+from utils.meta_data import MetaData
 
 _mimetypes = MimeTypes()
 _mimetypes.readfp(StringIO(mime_types))
@@ -73,28 +74,35 @@ def _get_file_type(file_id: str):
     return file_type
 
 
-def get_extension(file_id: str, mime_type: str) -> str:
+def get_extension(file_id: str, mime_type: str, dot: bool = True) -> str:
     """Get extension"""
 
+    if not file_id:
+        if dot:
+            return ".unknown"
+        return "unknown"
+
     file_type = _get_file_type(file_id)
 
     guessed_extension = _guess_extension(mime_type)
 
     if file_type in PHOTO_TYPES:
-        extension = ".jpg"
+        extension = "jpg"
     elif file_type == FileType.VOICE:
-        extension = guessed_extension or ".ogg"
+        extension = guessed_extension or "ogg"
     elif file_type in (FileType.VIDEO, FileType.ANIMATION, FileType.VIDEO_NOTE):
-        extension = guessed_extension or ".mp4"
+        extension = guessed_extension or "mp4"
     elif file_type == FileType.DOCUMENT:
-        extension = guessed_extension or ".zip"
+        extension = guessed_extension or "zip"
     elif file_type == FileType.STICKER:
-        extension = guessed_extension or ".webp"
+        extension = guessed_extension or "webp"
     elif file_type == FileType.AUDIO:
-        extension = guessed_extension or ".mp3"
+        extension = guessed_extension or "mp3"
     else:
-        extension = ".unknown"
+        extension = "unknown"
 
+    if dot:
+        extension = "." + extension
     return extension
 
 
@@ -549,3 +557,31 @@ async def check_user_permission(
         pass
 
     return False
+
+
+def set_meta_data(
+    meta_data: MetaData, message: pyrogram.types.Message, caption: str = None
+):
+    """Get all meta data"""
+    # message
+    meta_data.message_date = getattr(message, "date", None)
+    if caption:
+        meta_data.message_caption = caption
+    else:
+        meta_data.message_caption = getattr(message, "caption", None) or ""
+    meta_data.message_id = getattr(message, "id", None)
+    for kind in meta_data.AVAILABLE_MEDIA:
+        media_obj = getattr(message, kind, None)
+        if media_obj is not None:
+            meta_data.media_type = kind
+            break
+    else:
+        return
+    meta_data.media_file_name = getattr(media_obj, "file_name", None) or ""
+    meta_data.media_file_size = getattr(media_obj, "file_size", None)
+    meta_data.media_width = getattr(media_obj, "width", None)
+    meta_data.media_height = getattr(media_obj, "height", None)
+    meta_data.media_duration = getattr(media_obj, "duration", None)
+    meta_data.file_extension = get_extension(
+        media_obj.file_id, getattr(media_obj, "mime_type", ""), False
+    )
diff --git a/test_user_filter.py b/test_user_filter.py
@@ -1,6 +1,7 @@
 from datetime import datetime
 
 from module.filter import Filter, MetaData
+from module.pyrogram_extension import set_meta_data
 from tests.test_common import MockMessage, MockVideo
 from utils.format import replace_date_time
 
@@ -28,7 +29,7 @@
     ),
 )
 
-meta.get_meta_data(message)
+set_meta_data(meta, message)
 
 download_filter = Filter()
 

diff --git a/tests/test_common.py b/tests/test_common.py
@@ -1,6 +1,8 @@
 import datetime
 import platform
 
+from pyrogram.file_id import PHOTO_TYPES, FileType
+
 
 class Chat:
     def __init__(self, chat_id, chat_title):
@@ -123,3 +125,69 @@ def platform_generic_path(_path: str) -> str:
     if platform.system() == "Windows":
         platform_specific_path = platform_specific_path.replace("/", "\\")
     return platform_specific_path
+
+
+def get_file_type(file_id: str):
+    if file_id == "THUMBNAIL":
+        return FileType.THUMBNAIL
+    elif file_id == "CHAT_PHOTO":
+        return FileType.CHAT_PHOTO
+    elif file_id == "PHOTO":
+        return FileType.PHOTO
+    elif file_id == "VOICE":
+        return FileType.VOICE
+    elif file_id == "VIDEO":
+        return FileType.VIDEO
+    elif file_id == "DOCUMENT":
+        return FileType.DOCUMENT
+    elif file_id == "ENCRYPTED":
+        return FileType.ENCRYPTED
+    elif file_id == "TEMP":
+        return FileType.TEMP
+    elif file_id == "STICKER":
+        return FileType.STICKER
+    elif file_id == "AUDIO":
+        return FileType.AUDIO
+    elif file_id == "ANIMATION":
+        return FileType.ANIMATION
+    elif file_id == "ENCRYPTED_THUMBNAIL":
+        return FileType.ENCRYPTED_THUMBNAIL
+    elif file_id == "WALLPAPER":
+        return FileType.WALLPAPER
+    elif file_id == "VIDEO_NOTE":
+        return FileType.VIDEO_NOTE
+    elif file_id == "SECURE_RAW":
+        return FileType.SECURE_RAW
+    elif file_id == "SECURE":
+        return FileType.SECURE
+    elif file_id == "BACKGROUND":
+        return FileType.BACKGROUND
+    elif file_id == "DOCUMENT_AS_FILE":
+        return FileType.DOCUMENT_AS_FILE
+
+    raise ValueError("error file id!")
+
+
+def get_extension(file_id: str, mime_type: str, dot: bool = True):
+    file_type = get_file_type(file_id=file_id)
+    guessed_extension = ""
+
+    if file_type in PHOTO_TYPES:
+        extension = "jpg"
+    elif file_type == FileType.VOICE:
+        extension = guessed_extension or "ogg"
+    elif file_type in (FileType.VIDEO, FileType.ANIMATION, FileType.VIDEO_NOTE):
+        extension = guessed_extension or "mp4"
+    elif file_type == FileType.DOCUMENT:
+        extension = guessed_extension or "zip"
+    elif file_type == FileType.STICKER:
+        extension = guessed_extension or "webp"
+    elif file_type == FileType.AUDIO:
+        extension = guessed_extension or "mp3"
+    else:
+        extension = "unknown"
+
+    if dot:
+        extension = "." + extension
+
+    return extension
diff --git a/tests/test_media_downloader.py b/tests/test_media_downloader.py
@@ -10,7 +10,6 @@
 
 import mock
 import pyrogram
-from pyrogram.file_id import PHOTO_TYPES, FileType
 
 from media_downloader import (
     _can_download,
@@ -38,6 +37,7 @@
     MockVideo,
     MockVideoNote,
     MockVoice,
+    get_extension,
     platform_generic_path,
 )
 
@@ -164,69 +164,6 @@ def load_config():
     raise ValueError("error load config")
 
 
-def get_file_type(file_id: str):
-    if file_id == "THUMBNAIL":
-        return FileType.THUMBNAIL
-    elif file_id == "CHAT_PHOTO":
-        return FileType.CHAT_PHOTO
-    elif file_id == "PHOTO":
-        return FileType.PHOTO
-    elif file_id == "VOICE":
-        return FileType.VOICE
-    elif file_id == "VIDEO":
-        return FileType.VIDEO
-    elif file_id == "DOCUMENT":
-        return FileType.DOCUMENT
-    elif file_id == "ENCRYPTED":
-        return FileType.ENCRYPTED
-    elif file_id == "TEMP":
-        return FileType.TEMP
-    elif file_id == "STICKER":
-        return FileType.STICKER
-    elif file_id == "AUDIO":
-        return FileType.AUDIO
-    elif file_id == "ANIMATION":
-        return FileType.ANIMATION
-    elif file_id == "ENCRYPTED_THUMBNAIL":
-        return FileType.ENCRYPTED_THUMBNAIL
-    elif file_id == "WALLPAPER":
-        return FileType.WALLPAPER
-    elif file_id == "VIDEO_NOTE":
-        return FileType.VIDEO_NOTE
-    elif file_id == "SECURE_RAW":
-        return FileType.SECURE_RAW
-    elif file_id == "SECURE":
-        return FileType.SECURE
-    elif file_id == "BACKGROUND":
-        return FileType.BACKGROUND
-    elif file_id == "DOCUMENT_AS_FILE":
-        return FileType.DOCUMENT_AS_FILE
-
-    raise ValueError("error file id!")
-
-
-def get_extension(file_id: str, mime_type: str):
-    file_type = get_file_type(file_id=file_id)
-    guessed_extension = ""
-
-    if file_type in PHOTO_TYPES:
-        extension = ".jpg"
-    elif file_type == FileType.VOICE:
-        extension = guessed_extension or ".ogg"
-    elif file_type in (FileType.VIDEO, FileType.ANIMATION, FileType.VIDEO_NOTE):
-        extension = guessed_extension or ".mp4"
-    elif file_type == FileType.DOCUMENT:
-        extension = guessed_extension or ".zip"
-    elif file_type == FileType.STICKER:
-        extension = guessed_extension or ".webp"
-    elif file_type == FileType.AUDIO:
-        extension = guessed_extension or ".mp3"
-    else:
-        extension = ".unknown"
-
-    return extension
-
-
 class MyQueue:
     def __init__(self, queue_list_obj):
         self._queue = queue.Queue()