getstickers 支持自定义 emoji 包

getstickers/main.py

@@ -4,10 +4,11 @@ from collections import defaultdict
 import os
 import zipfile
 
+from pyrogram.enums import MessageEntityType
 from pyrogram.raw.functions.messages import GetStickerSet
 from pyrogram.raw.types import InputStickerSetShortName
 from pyrogram.raw.types.messages import StickerSet
-from pyrogram.types import Document
+from pyrogram.types import Document, Sticker
 
 from pagermaid import working_dir
 from pagermaid.enums import Message, Client
@@ -15,24 +16,13 @@ from pagermaid.listener import listener
 from pagermaid.single_utils import safe_remove
 
 
-@listener(command="getstickers",
-          description="获取整个贴纸包的贴纸")
-async def get_stickers(bot: Client, message: Message):
-    if not os.path.isdir('data/sticker/'):
-        os.makedirs('data/sticker/')
-    if message.reply_to_message:
-        if not message.reply_to_message.sticker:
-            return await message.edit("请回复一张贴纸。")
-    else:
-        return await message.edit("请回复一张贴纸。")
-    sticker = message.reply_to_message.sticker
-    if not sticker.set_name:
-        return await message.edit("回复的贴纸不属于任何贴纸包。")
+async def download_stickers(bot: Client, message: Message, sticker: Sticker):
     try:
         sticker_set: StickerSet = await bot.invoke(
             GetStickerSet(stickerset=InputStickerSetShortName(short_name=sticker.set_name), hash=0))
-    except Exception:
+    except Exception:  # noqa
         return await message.edit('回复的贴纸不存在于任何贴纸包中。')
+
     pack_file = os.path.join('data/sticker/', sticker_set.set.short_name, "pack.txt")
 
     if os.path.isfile(pack_file):
@@ -58,13 +48,17 @@ async def get_stickers(bot: Client, message: Message):
         download(document, emojis, f"data/sticker/{sticker_set.set.short_name}", f"{i:03d}.{file_ext_ns_ion}")
     ) for i, document in enumerate(sticker_set.documents)]
 
-    message: Message = await message.edit(f"正在下载 {sticker_set.set.short_name} 中的 {sticker_set.set.count} 张贴纸。。。")
+    message: Message = await message.edit(
+        f"正在下载 {sticker_set.set.short_name} 中的 {sticker_set.set.count} 张贴纸。。。")
 
     while 1:
         done, pending_tasks = await asyncio.wait(pending_tasks, timeout=2.5, return_when=asyncio.FIRST_COMPLETED)
         if not pending_tasks:
             break
+    await upload_sticker(bot, message, sticker_set)
 
+
+async def upload_sticker(bot: Client, message: Message, sticker_set: StickerSet):
     await message.edit("下载完毕，打包上传中。")
     directory_name = sticker_set.set.short_name
     os.chdir("data/sticker/")
@@ -83,6 +77,33 @@ async def get_stickers(bot: Client, message: Message):
     await message.safe_delete()
 
 
+async def get_custom_emojis(bot: Client, message: Message):
+    if message.entities:
+        for entity in message.entities:
+            if entity.type == MessageEntityType.CUSTOM_EMOJI:
+                try:
+                    sticker = await bot.get_custom_emoji_stickers([entity.custom_emoji_id])
+                except Exception:
+                    return None
+                return sticker[0] if sticker else None
+
+
+@listener(command="getstickers",
+          description="获取整个贴纸包的贴纸")
+async def get_stickers(bot: Client, message: Message):
+    if not os.path.isdir('data/sticker/'):
+        os.makedirs('data/sticker/')
+    if message.reply_to_message:
+        sticker = message.reply_to_message.sticker or await get_custom_emojis(bot, message.reply_to_message)
+    else:
+        sticker = message.sticker or await get_custom_emojis(bot, message)
+    if not sticker:
+        return await message.edit("请回复一张贴纸。")
+    if not sticker.set_name:
+        return await message.edit("回复的贴纸不属于任何贴纸包。")
+    await download_stickers(bot, message, sticker)
+
+
 def zipdir(path, zip_):
     for root, dirs, files in os.walk(path):
         for file in files: