Ver código fonte

Revert "Rewrite of the kang (stickers) module. (#50)" (#54)

This reverts commit 211c704530b75014ad4750000f980a81f3fc3d05.
Akshay Rajput 5 anos atrás
pai
commit
eb93279f04
3 arquivos alterados com 112 adições e 129 exclusões
  1. 112 48
      wbb/modules/stickers.py
  2. 0 43
      wbb/utils/files.py
  3. 0 38
      wbb/utils/stickerset.py

+ 112 - 48
wbb/modules/stickers.py

@@ -1,14 +1,14 @@
 import os
 import math
-import uuid
 import imghdr
+from PIL import Image
+from telegram import Bot as tg
+from telegram import TelegramError
 from wbb.utils.botinfo import BOT_USERNAME
 from wbb import BOT_TOKEN
 from wbb import app
 from pyrogram import filters
 from wbb.utils.errors import capture_err
-from wbb.utils.files import resize_file_to_sticker_size, upload_document, get_document_from_file_id
-from wbb.utils.stickerset import get_sticker_set_by_name, create_sticker, add_sticker_to_set, create_sticker_set
 from random import randint
 from pyrogram.types import InlineKeyboardButton, InlineKeyboardMarkup
 
@@ -17,8 +17,8 @@ __MODULE__ = "Stickers"
 __HELP__ = """/sticker_id - To Get File ID of A Sticker.
 /kang - To Kang A Sticker or Image."""
 
-MAX_STICKERS = 120 # would be better if we could fetch this limit directly from telegram
-SUPPORTED_TYPES = ['jpeg', 'png', 'webp']
+## Another marie based kang module
+
 
 @app.on_message(filters.command("sticker_id") & ~filters.edited)
 @capture_err
@@ -33,59 +33,123 @@ async def sticker_id(_, message):
     await message.reply_text(f"`{file_id}`")
 
 
+updater = tg(BOT_TOKEN)
+
+
 @app.on_message(filters.command("kang") & ~filters.edited)
 @capture_err
 async def kang(client, message):
     if not message.reply_to_message:
         await message.reply_text("Reply to a sticker/image to kang it.")
         return
-    msg = await message.reply_text("Kanging Sticker..")
-
-    # Find the proper emoji
+    kangsticker = f"./{randint(10000, 99999)}.png"
     args = message.text.split()
-    if len(args) > 1:
-        sticker_emoji = str(args[1])
-    elif message.reply_to_message.sticker and message.reply_to_message.sticker.emoji:
-        sticker_emoji = message.reply_to_message.sticker.emoji
-    else:
-        sticker_emoji = "🤔"
-
-    # Get the corresponding fileid, resize the file if necessary
-    file_id = (message.reply_to_message.photo or message.reply_to_message.document)
-    if message.reply_to_message.sticker:
-        sticker = await create_sticker(await get_document_from_file_id(message.reply_to_message.sticker.file_id), sticker_emoji)
-    elif file_id:
-        temp_file_path = await app.download_media(file_id, file_name=str(uuid.uuid4()))
-        if imghdr.what(temp_file_path) not in SUPPORTED_TYPES:
+    msg = await message.reply_text("Kanging Sticker..")
+    user = message.from_user
+    packnum = 0
+    packname = "f" + str(user.id) + "_by_"+ BOT_USERNAME
+    packname_found = 0
+    max_stickers = 120
+    while packname_found == 0:
+        try:
+            stickerset = updater.get_sticker_set(packname)
+            if len(stickerset.stickers) >= max_stickers:
+                packnum += 1
+                packname = "f" + str(packnum) + "_" + \
+                    str(user.id) + "_by_"+BOT_USERNAME
+            else:
+                packname_found = 1
+        except TelegramError as e:
+            if e.message == "Stickerset_invalid":
+                packname_found = 1
+    if message.reply_to_message:
+        if message.reply_to_message.sticker:
+            file_id = message.reply_to_message.sticker.file_id
+        elif message.reply_to_message.photo:
+            file_id = message.reply_to_message.photo.file_id
+        elif message.reply_to_message.document:
+            file_id = message.reply_to_message.document.file_id
+        else:
+            await msg.edit("Nope, can't kang that.")
+            return
+        await app.download_media(file_id, file_name=kangsticker)
+        image_type = imghdr.what(kangsticker)
+        if image_type != 'jpeg' and image_type != 'png' and image_type != 'webp':
             await msg.edit("Format not supported! ({})".format(image_type))
             return
+        if len(args) > 1:
+            sticker_emoji = str(args[1])
+        elif message.reply_to_message.sticker and message.reply_to_message.sticker.emoji:
+            sticker_emoji = message.reply_to_message.sticker.emoji
+        else:
+            sticker_emoji = "🤔"
         try:
-            await resize_file_to_sticker_size(temp_file_path)
+            im = Image.open(kangsticker)
+            maxsize = (512, 512)
+            if (im.width and im.height) < 512:
+                size1 = im.width
+                size2 = im.height
+                if im.width > im.height:
+                    scale = 512/size1
+                    size1new = 512
+                    size2new = size2 * scale
+                else:
+                    scale = 512/size2
+                    size1new = size1 * scale
+                    size2new = 512
+                size1new = math.floor(size1new)
+                size2new = math.floor(size2new)
+                sizenew = (size1new, size2new)
+                im = im.resize(sizenew)
+            else:
+                im.thumbnail(maxsize)
+            if not message.reply_to_message.sticker:
+                im.save(kangsticker, "PNG")
+            updater.add_sticker_to_set(user_id=user.id, name=packname,
+                                       png_sticker=open(kangsticker, 'rb'), emojis=sticker_emoji)
+            await msg.edit("Sticker Kanged To [Pack](t.me/addstickers/{})\nEmoji: {}".format(packname, sticker_emoji))
         except OSError as e:
-            await msg.edit_text("Something wrong happened.")
-            raise Exception(f"Something went wrong while resizing the sticker (at {temp_file_path}); {e}")
-            return False
-        sticker =  await create_sticker(await upload_document(client, temp_file_path), sticker_emoji)
-        if os.path.isfile(temp_file_path):
-            os.remove(temp_file_path)
-    else:
-        await msg.edit("Nope, can't kang that.")
-        return
+            await message.reply_text("Something wrong happened.")
+            print(e)
+            return
+        except TelegramError as e:
+            if e.message == "Stickerset_invalid":
+                await makepack_internal(msg, user, open(kangsticker, 'rb'), sticker_emoji, updater, packname, packnum)
+            elif e.message == "Sticker_png_dimensions":
+                im.save(kangsticker, "PNG")
+                updater.add_sticker_to_set(user_id=user.id, name=packname,
+                                           png_sticker=open(kangsticker, 'rb'), emojis=sticker_emoji)
+                await msg.edit("Sticker Kanged [pack](t.me/addstickers/{})\nEmoji: ```{}```".format(packname, sticker_emoji))
+            elif e.message == "Invalid sticker emojis":
+                await msg.edit("Invalid emoji")
+            elif e.message == "Stickers_too_much":
+                await msg.edit("Too many stickers")
+            elif e.message == "Internal Server Error: sticker set not found (500)":
+                await msg.edit("Sticker Kanged [pack](t.me/addstickers/{})\nEmoji: ```{}```".format(packname, sticker_emoji))
+    if os.path.isfile(kangsticker):
+        os.remove(kangsticker)
 
-    # Find an available pack & add the sticker to the pack; create a new pack if needed
-    packnum = 0 # Would be a good idea to cache the number instead of searching it every single time...
-    packname = "f" + str(message.from_user.id) + "_by_"+ BOT_USERNAME
-    while True:
-        stickerset = await get_sticker_set_by_name(client, packname)
-        if not stickerset:
-            stickerset = await create_sticker_set(client, message.from_user.id, f"{message.from_user.first_name[:32]}'s kang pack", packname, [sticker])
-        elif stickerset.set.count >= MAX_STICKERS:
-            packnum += 1
-            packname = "f" + str(packnum) + "_" + \
-                str(message.from_user.id) + "_by_"+BOT_USERNAME
-            continue
-        else:
-            add_sticker_to_set(client, stickerset, sticker)
-        break
 
-    await msg.edit("Sticker Kanged To [Pack](t.me/addstickers/{})\nEmoji: {}".format(packname, sticker_emoji))
+async def makepack_internal(msg, user, png_sticker, emoji, updater, packname, packnum):
+    name = user.first_name
+    name = name[:50]
+    try:
+        extra_version = ""
+        if packnum > 0:
+            extra_version = " " + str(packnum)
+        success = updater.create_new_sticker_set(user.id, packname, f"{name}'s wbb pack" + extra_version,
+                                                 png_sticker=png_sticker,
+                                                 emojis=emoji)
+    except TelegramError as e:
+        if e.message == "Sticker set name is already occupied":
+            await msg.edit("Your pack can be found [here](t.me/addstickers/%s)" % packname)
+        elif e.message == "Peer_id_invalid":
+            await msg.edit("Contact me in PM first.", reply_markup=InlineKeyboardMarkup([[InlineKeyboardButton(
+                text="Start", url=f"t.me/{BOT_USERNAME}")]]))
+        elif e.message == "Internal Server Error: created sticker set not found (500)":
+            await msg.edit("Sticker pack successfully created. Get it [here](t.me/addstickers/%s)" % packname)
+        return
+    if success:
+        await msg.edit("Sticker pack successfully created. Get it [here](t.me/addstickers/%s)" % packname)
+    else:
+        await msg.edit("Failed to create sticker pack.")

+ 0 - 43
wbb/utils/files.py

@@ -1,43 +0,0 @@
-from PIL import Image
-
-from pyrogram import Client, raw
-from pyrogram.file_id import FileId
-
-STICKER_DIMENSIONS = (512, 512)
-
-
-async def resize_file_to_sticker_size(file_path: str):
-    im = Image.open(file_path)
-    if (im.width, im.height) < STICKER_DIMENSIONS:
-        size1 = im.width
-        size2 = im.height
-        if im.width > im.height:
-            scale = STICKER_DIMENSIONS[0]/size1
-            size1new = STICKER_DIMENSIONS[0]
-            size2new = size2 * scale
-        else:
-            scale = STICKER_DIMENSIONS[1]/size2
-            size1new = size1 * scale
-            size2new = STICKER_DIMENSIONS[1]
-        size1new = math.floor(size1new)
-        size2new = math.floor(size2new)
-        sizenew = (size1new, size2new)
-        im = im.resize(sizenew)
-    else:
-        im.thumbnail(STICKER_DIMENSIONS)
-    im.save(file_path, "PNG")
-
-async def upload_document(client: Client, file_path: str) -> raw.base.InputDocument:
-    media = await client.send(
-        raw.functions.messages.UploadMedia(
-            peer=raw.types.InputPeerEmpty(),
-            media=raw.types.InputMediaUploadedPhoto(
-                file=file_path
-            )
-        )
-    )
-    return raw.types.InputDocument(id=media.photo.id, access_hash=media.access_hash, file_reference=media.photo.file_reference)
-
-async def get_document_from_file_id(file_id: str) -> raw.base.InputDocument:
-    decoded = FileId.decode(file_id)
-    return raw.types.InputDocument(id=decoded.media_id, access_hash=decoded.access_hash, file_reference=decoded.file_reference)

+ 0 - 38
wbb/utils/stickerset.py

@@ -1,38 +0,0 @@
-from typing import List
-
-from pyrogram import Client, raw
-from pyrogram import errors
-
-async def get_sticker_set_by_name(client: Client, name: str) -> raw.base.messages.StickerSet:
-    try:
-        return await client.send(
-            raw.functions.messages.GetStickerSet(
-                stickerset = raw.types.InputStickerSetShortName(short_name=name)
-            )
-        )
-    except errors.exceptions.not_acceptable_406.StickersetInvalid:
-        return None
-
-# Known errors: (I don't see a reason to catch them as we, for sure, won't face them right now):
-# errors.exceptions.bad_request_400.PackShortNameInvalid -> pack name needs to end with _by_botname
-# errors.exceptions.bad_request_400.ShortnameOccupyFailed -> pack's name is already in use
-async def create_sticker_set(client: Client, owner: int, title: str, short_name: str, stickers: List[raw.base.InputStickerSetItem]) -> raw.base.messages.StickerSet:
-    return await client.send(
-        raw.functions.stickers.CreateStickerSet(
-            user_id=await client.resolve_peer(owner), 
-            title=title,
-            short_name=short_name,
-            stickers=stickers
-        )
-    )
-
-async def add_sticker_to_set(client: Client, stickerset: raw.base.messages.StickerSet, sticker: raw.base.InputStickerSetItem) -> raw.base.messages.StickerSet:
-    return await client.send(
-        raw.functions.stickers.AddStickerToSet(
-            stickerset=stickerset,
-            sticker=sticker
-        )
-    )
-
-async def create_sticker(sticker: raw.base.InputDocument, emoji: str) -> raw.base.InputStickerSetItem:
-    return raw.types.InputStickerSetItem(document=sticker, emoji=emoji)