Commit f4571be

benny-dou <60535774+benny-dou@users.noreply.github.com>
2025-03-20 16:29:25
feat: add `/convert` command to convert image file to photo
1 parent 813f6d2
src/others/raw_img_file.py
@@ -6,26 +6,74 @@ from loguru import logger
 from pyrogram.client import Client
 from pyrogram.types import Message
 
+from config import PREFIX
 from messages.parser import parse_msg
 from messages.sender import send2tg
+from messages.utils import equal_prefix, startswith_prefix
 from multimedia import split_long_img
 
+HELP = f"""🔁**原图转图片**
+使用说明: 以 `{PREFIX.CONVERT}` 回复包含原图文件的消息即可
+"""
 
-async def convert_raw_img_file(client: Client, message: Message, **kwargs):
-    mime_type = message.document.mime_type if message.document else ""
-    if mime_type not in ["image/png", "image/jpeg", "image/heic"]:
+
+async def convert_raw_img_file(client: Client, message: Message, *, convert_need_prefix: bool | None = None, **kwargs):
+    # send docs if message == "/convert", without reply
+    if equal_prefix(message.text, prefix=[PREFIX.CONVERT]) and not message.reply_to_message:
+        await send2tg(client, message, texts=HELP, **kwargs)
+        return
+    if not (trigger_message := get_trigger_message(message, convert_need_prefix)):
+        return
+    info = parse_msg(trigger_message)
+    if info["mime_type"] not in ["image/png", "image/jpeg", "image/heic"]:
         return
     try:
-        info = parse_msg(message)
-        logger.info(f"Splitting raw image: {info['file_name']}")
-        path: str = await message.download()  # type: ignore
-        logger.trace(f"Image {{info['file_name']}} downloaded to {path}")
+        logger.info(f"Convert raw image: {info['file_name']}")
+        path: str = await trigger_message.download()  # type: ignore
+        logger.trace(f"Image {info['file_name']} downloaded to {path}")
         photos = split_long_img(path)
         media = [{"photo": photo.as_posix()} for photo in photos if photo.is_file()]
         # send splits
         if len(media) == 1:
-            await send2tg(client, message, media=media, texts="🔁【原图转换】已完成", **kwargs)
+            await send2tg(client, trigger_message, media=media, texts="🔁【原图转换】已完成", **kwargs)
         else:
-            await send2tg(client, message, media=media, texts="✂️【自动切长图】已完成 (使用【原图】发送才可触发切长图)", **kwargs)
+            await send2tg(client, trigger_message, media=media, texts="✂️【自动切长图】已完成 (使用【原图】发送才可触发切长图)", **kwargs)
     except Exception as e:
         logger.exception(e)
+
+
+def get_trigger_message(
+    message: Message,
+    convert_need_prefix: bool | None = None,
+) -> Message | None:
+    """Check if the message is triggerable for converting.
+
+    By default, "/convert" prefix is needed in in Group & Channel & Bot chats to trigger this function.
+    In private chat, no need to add "/convert" prefix
+    """
+    info = parse_msg(message)
+    this_text = info["text"]  # this message
+    if info["ctype"].lower() in ["group", "supergroup", "channel", "bot"]:  # noqa: SIM108
+        convert_need_prefix = convert_need_prefix if convert_need_prefix is not None else True
+    else:  # private chat
+        convert_need_prefix = convert_need_prefix or False
+
+    # skip no "/convert" prefix message if convert_need_prefix
+    if convert_need_prefix and not startswith_prefix(this_text, prefix=[PREFIX.CONVERT]):
+        return None
+
+    # treat the reply_to_message as the real message need to be converted
+    trigger_msg = message.reply_to_message if message.reply_to_message and startswith_prefix(this_text, prefix=[PREFIX.CONVERT]) else message
+    trigger_info = parse_msg(trigger_msg, silent=True)
+
+    # skip non document message
+    if not trigger_msg:
+        return None
+    if trigger_info["mtype"] not in ["document"]:
+        return None
+
+    # always trigger if the message has "/convert" prefix
+    if startswith_prefix(this_text, prefix=[PREFIX.CONVERT]):
+        return trigger_msg
+
+    return trigger_msg
src/config.py
@@ -62,6 +62,7 @@ class PREFIX:
     AI_SUMMARY = os.getenv("PREFIX_AI_SUMMARY", "/summary").lower()
     ASR = os.getenv("PREFIX_ASR", "/asr").lower()
     AUDIO = os.getenv("PREFIX_AUDIO", "/audio").lower()
+    CONVERT = os.getenv("PREFIX_CONVERT", "/convert").lower()  # convert image file to photo
     GPT = os.getenv("PREFIX_GPT", "/ai").lower()
     SUBTITLE = os.getenv("PREFIX_SUBTITLE", "/subtitle").lower()
     WGET = os.getenv("PREFIX_WGET", "/wget").lower()
src/handler.py
@@ -48,7 +48,7 @@ async def handle_utilities(
     ocr: bool = True,
     price: bool = True,
     summary: bool = True,
-    raw_img: bool = False,
+    raw_img: bool = True,
     show_progress: bool = True,
     detail_progress: bool = False,
     **kwargs,
src/permission.py
@@ -91,7 +91,6 @@ def check_service(cid: int | str, ctype: str) -> dict:
 
     permission = {
         # default to False
-        "raw_img": False,  # only enable in private chat
         "ai": False,
         "prepend_sender_user": False,
         # default to True
@@ -103,6 +102,7 @@ def check_service(cid: int | str, ctype: str) -> dict:
         "wget": True,
         "ocr": True,
         "price": True,
+        "raw_img": True,
         "summary": True,
         "show_progress": True,
         "detail_progress": True,
@@ -117,7 +117,6 @@ def check_service(cid: int | str, ctype: str) -> dict:
 
     if ctype == "PRIVATE":
         permission["ai"] = True
-        permission["raw_img"] = True
         permission["need_prefix"] = False
 
     # global service permission