Rewrite of the kang module (v2) (#55)

* Rewrite of the kang (stickers) module.
I changed the module to use the pyrogram's built-in raw functions

> I fixed a lot of problems and ended up rewriting it

* Small fixes to the kang module

* More fixes in the kang module

* Small fix regarding file formats in the kang module

wbb/modules/stickers.py

@@ -1,15 +1,12 @@
 import os
-import math
 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 random import randint
+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 pyrogram.types import InlineKeyboardButton, InlineKeyboardMarkup
 
 
@@ -17,8 +14,8 @@ __MODULE__ = "Stickers"
 __HELP__ = """/sticker_id - To Get File ID of A Sticker.
 /kang - To Kang A Sticker or Image."""
 
-## Another marie based kang module
-
+MAX_STICKERS = 120 # would be better if we could fetch this limit directly from telegram
+SUPPORTED_TYPES = ['jpeg', 'png', 'webp']
 
 @app.on_message(filters.command("sticker_id") & ~filters.edited)
 @capture_err
@@ -33,123 +30,59 @@ 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
-    kangsticker = f"./{randint(10000, 99999)}.png"
-    args = message.text.split()
     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':
+
+    # Find the proper emoji
+    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
+    doc = (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 doc:
+        temp_file_path = await app.download_media(doc)
+        if imghdr.what(temp_file_path) not in SUPPORTED_TYPES:
             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:
-            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))
+            temp_file_path = await resize_file_to_sticker_size(temp_file_path)
         except OSError as e:
-            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)
+            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, message.chat.id), sticker_emoji)
+        if os.path.isfile(temp_file_path):
+            os.remove(temp_file_path)
+    else:
+        await msg.edit("Nope, can't kang that.")
+        return
 
+    # 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:
+            await add_sticker_to_set(client, stickerset, sticker)
+        break
 
-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.")
+    await msg.edit("Sticker Kanged To [Pack](t.me/addstickers/{})\nEmoji: {}".format(packname, sticker_emoji))

wbb/utils/files.py

@@ -0,0 +1,56 @@
+import os
+from PIL import Image
+import math
+
+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) -> 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)
+    try:
+        os.remove(file_path)
+        file_path = f"{file_path}.png"
+        return file_path
+    finally:
+        im.save(file_path)
+
+async def upload_document(client: Client, file_path: str, chat_id: int) -> raw.base.InputDocument:
+    media = await client.send(
+        raw.functions.messages.UploadMedia(
+            peer=await client.resolve_peer(chat_id),
+            media=raw.types.InputMediaUploadedDocument(
+                mime_type=client.guess_mime_type(file_path) or "application/zip",
+                file=await client.save_file(file_path),
+                attributes=[
+                    raw.types.DocumentAttributeFilename(
+                        file_name=os.path.basename(file_path)
+                    )
+                ]
+            )
+        )
+    )
+    return raw.types.InputDocument(id=media.document.id, access_hash=media.document.access_hash, file_reference=media.document.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)

wbb/utils/stickerset.py

@@ -0,0 +1,38 @@
+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=raw.types.InputStickerSetShortName(short_name=stickerset.set.short_name),
+            sticker=sticker
+        )
+    )
+
+async def create_sticker(sticker: raw.base.InputDocument, emoji: str) -> raw.base.InputStickerSetItem:
+    return raw.types.InputStickerSetItem(document=sticker, emoji=emoji)