diff --git a/getstickers/main.py b/getstickers/main.py index 6a98deb..32b265b 100644 --- a/getstickers/main.py +++ b/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: