Pārlūkot izejas kodu

added: support to all kinda notes (#225)

SI_NN_ER_LS 2 gadi atpakaļ
vecāks
revīzija
5b72f233eb
3 mainītis faili ar 147 papildinājumiem un 44 dzēšanām
  1. 19 19
      wbb/modules/greetings.py
  2. 121 19
      wbb/modules/notes.py
  3. 7 6
      wbb/utils/dbfunctions.py

+ 19 - 19
wbb/modules/greetings.py

@@ -226,7 +226,7 @@ async def welcome(_, message: Message):
         return await handle_new_member(message, member, chat)
 
 async def send_welcome_message(chat: Chat, user_id: int, delete: bool = False):
-    raw_text, animation_id, photo_id = await get_welcome(chat.id)
+    welcome, raw_text, file_id = await get_welcome(chat.id)
 
     if not raw_text:
         return
@@ -238,24 +238,24 @@ async def send_welcome_message(chat: Chat, user_id: int, delete: bool = False):
         text = text.replace("{name}", (await app.get_users(user_id)).mention)
 
     async def _send_wait_delete():
-        if not animation_id and not photo_id:
+        if welcome == "Text":
             m = await app.send_message(
                 chat.id,
                 text=text,
                 reply_markup=keyb,
                 disable_web_page_preview=True,
             )
-        elif not animation_id:
+        elif welcome == "Photo":
             m = await app.send_photo(
                 chat.id,
-                photo=photo_id,
+                photo=file_id,
                 caption=text,
                 reply_markup=keyb,
             )
         else:
             m = await app.send_animation(
                 chat.id,
-                animation=animation_id,
+                animation=file_id,
                 caption=text,
                 reply_markup=keyb,
             )
@@ -397,11 +397,11 @@ async def captcha_state(_, message):
 # WELCOME MESSAGE
 
 
-async def check_caption(message, chat_id, raw_text, animation_id, photo_id):
+async def check_caption(message, chat_id, welcome, raw_text, file_id):
     if not extract_text_and_keyb(ikb, raw_text):
         return await message.reply_text("Wrong formatting, check the help section.")
     
-    await set_welcome(chat_id, raw_text, animation_id, photo_id)
+    await set_welcome(chat_id, welcome, raw_text, file_id)
     await message.reply_text("Welcome message has been successfully set.")
 
 @app.on_message(filters.command("set_welcome") & ~filters.private)
@@ -425,27 +425,27 @@ async def set_welcome_func(_, message):
         return await message.reply_text(usage, reply_markup=key)
 
     if replied_message.animation:
-        animation_id = replied_message.animation.file_id
-        text = replied_message.caption
-        photo_id = None
+        welcome = "Animation"
+        file_id = replied_message.animation.file_id
+        text = replied_message.caption      
         if not text:
             return await message.reply_text(usage, reply_markup=key)
         raw_text = text.markdown
-        return await check_caption(message, chat_id, raw_text, animation_id, photo_id)
+        return await check_caption(message, chat_id, welcome, raw_text, file_id)
     if replied_message.photo:
-        photo_id = replied_message.photo.file_id
+        welcome = "Photo"
+        file_id = replied_message.photo.file_id
         text = replied_message.caption
-        animation_id = None
         if not text:
             return await message.reply_text(usage, reply_markup=key)
         raw_text = text.markdown
-        return await check_caption(message, chat_id, raw_text, animation_id, photo_id)
+        return await check_caption(message, chat_id, welcome, raw_text, file_id)
     if replied_message.text:
-        animation_id = None
+        welcome = "Text"
+        file_id = None
         text = replied_message.text
-        photo_id = None
         raw_text = text.markdown
-        return await check_caption(message, chat_id, raw_text, animation_id, photo_id)
+        return await check_caption(message, chat_id, welcome, raw_text, file_id)
     else:
         await message.reply_text("Only text, gif and photo welcome message are supposed")
 
@@ -462,7 +462,7 @@ async def del_welcome_func(_, message):
 @adminsOnly("can_change_info")
 async def get_welcome_func(_, message):
     chat = message.chat
-    raw_text, animation_id, photo_id = await get_welcome(chat.id)
+    welcome, raw_text, file_id = await get_welcome(chat.id)
     if not raw_text:
         return await message.reply_text("No welcome message set.")
     if not message.from_user:
@@ -470,4 +470,4 @@ async def get_welcome_func(_, message):
 
     await send_welcome_message(chat, message.from_user.id)
 
-    await message.reply_text(f'photo_id: `{photo_id}`\n\nGif_id: `{animation_id}`\n\n`{raw_text.replace("`", "")}`')
+    await message.reply_text(f'Welcome: {welcome}\n\nFile_id: `{file_id}`\n\n`{raw_text.replace("`", "")}`')

+ 121 - 19
wbb/modules/notes.py

@@ -35,7 +35,11 @@ from wbb.utils.functions import extract_text_and_keyb
 __MODULE__ = "Notes"
 __HELP__ = """/notes To Get All The Notes In The Chat.
 
-/save [NOTE_NAME] To Save A Note (Can be a sticker or text).
+/save [NOTE_NAME] To Save A Note.
+
+Supported note types are Text, Animation, Photo, Document, Video, video notes, Audio, Voice
+
+To change caption of any files use.\n/save [NOTE_NAME] [NEW_CAPTION]
 
 #NOTE_NAME To Get A Note.
 
@@ -45,6 +49,17 @@ Checkout /markdownhelp to know more about formattings and other syntax.
 """
 
 
+def extract_urls(reply_markup):
+    urls = []
+    if reply_markup.inline_keyboard:
+        buttons = reply_markup.inline_keyboard
+        for i, row in enumerate(buttons):
+            for j, button in enumerate(row):
+                 if button.url:
+                    name = "\n~\nbutton" if i * len(row) + j + 1 == 1 else f"button{i * len(row) + j + 1}"
+                    urls.append((f"{name}", button.text, button.url))
+    return urls
+
 @app2.on_message(
     filters.command("save", prefixes=USERBOT_PREFIX)
     & ~filters.forwarded
@@ -57,24 +72,63 @@ async def save_notee(_, message):
     if len(message.command) < 2 or not message.reply_to_message:
         await eor(
             message,
-            text="**Usage:**\nReply to a text or sticker with /save [NOTE_NAME] to save it.",
-        )
-
-    elif not message.reply_to_message.text and not message.reply_to_message.sticker:
-        await eor(
-            message,
-            text="__**You can only save text or stickers in notes.**__",
+            text="**Usage:**\nReply to a message with /save [NOTE_NAME] to save a new note.",
         )
     else:
         name = message.text.split(None, 1)[1].strip()
         if not name:
             return await eor(message, text="**Usage**\n__/save [NOTE_NAME]__")
-        _type = "text" if message.reply_to_message.text else "sticker"
+        replied_message = message.reply_to_message
+        text = name.split(" ", 1)
+        if len(text) > 1:
+            name = text[0]
+            data = text[1].strip()
+            if replied_message.sticker or replied_message.video_note:
+                data = None
+        else:
+            if replied_message.sticker or replied_message.video_note:
+                data = None
+            elif not replied_message.text and not replied_message.caption:
+                data = None
+            else:
+                data = replied_message.text.markdown if replied_message.text else replied_message.caption.markdown
+        if replied_message.text:
+            _type = "text"
+            file_id = None
+        if replied_message.sticker:
+            _type = "sticker"
+            file_id = replied_message.sticker.file_id
+        if replied_message.animation:
+            _type = "animation"
+            file_id = replied_message.animation.file_id
+        if replied_message.photo:
+            _type = "photo"
+            file_id = replied_message.photo.file_id
+        if replied_message.document:
+            _type = "document"
+            file_id = replied_message.document.file_id
+        if replied_message.video:
+            _type = "video"
+            file_id = replied_message.video.file_id
+        if replied_message.video_note:
+            _type = "video_note"
+            file_id = replied_message.video_note.file_id
+        if replied_message.audio:
+            _type = "audio"
+            file_id = replied_message.audio.file_id
+        if replied_message.voice:
+            _type = "voice"
+            file_id = replied_message.voice.file_id
+        if replied_message.reply_markup:
+            if not "~" in data:
+                urls = extract_urls(replied_message.reply_markup)
+                if urls:
+                    response = "\n".join([f"{name}=[{text}, {url}]" for name, text, url in urls])
+                    data = data + response
         note = {
             "type": _type,
-            "data": message.reply_to_message.text.markdown
-            if _type == "text"
-            else message.reply_to_message.sticker.file_id,
+            "data": data,
+            "file_id": file_id,
         }
         prefix = message.text.split()[0][0]
         chat_id = message.chat.id if prefix != USERBOT_PREFIX else USERBOT_ID
@@ -142,21 +196,69 @@ async def get_one_note(_, message):
     _note = await get_note(message.chat.id, name)
     if not _note:
         return
-    if _note["type"] == "text":
-        data = _note["data"]
-        keyb = None
+    type = _note["type"]
+    data = _note["data"]
+    file_id = _note["file_id"]
+    keyb = None
+    if data:       
         if findall(r"\[.+\,.+\]", data):
             keyboard = extract_text_and_keyb(ikb, data)
             if keyboard:
                 data, keyb = keyboard
+    replied_message = message.reply_to_message
+    if replied_message:
+        if replied_message.from_user.id != message.from_user.id:
+            message = replied_message
+    if type == "text":
         await message.reply_text(
-            data,
+            text=data,
             reply_markup=keyb,
             disable_web_page_preview=True,
         )
-    else:
-        await message.reply_sticker(_note["data"])
-
+    if type == "sticker":
+        await message.reply_sticker(
+            sticker=file_id,
+        )
+    if type == "animation":
+        await message.reply_animation(
+            animation=file_id,
+            caption=data,
+            reply_markup=keyb,
+        )
+    if type == "photo":
+        await message.reply_photo(
+            photo=file_id,
+            caption=data,
+            reply_markup=keyb,
+        )
+    if type == "document":
+        await message.reply_document(
+            document=file_id,
+            caption=data,
+            reply_markup=keyb,
+        )
+    if type == "video":
+        await message.reply_video(
+            video=file_id,
+            caption=data,
+            reply_markup=keyb,
+        )
+    if type == "video_note":
+        await message.reply_video_note(
+            video_note=file_id,
+        )
+    if type == "audio":
+        await message.reply_audio(
+            audio=file_id,
+            caption=data,
+            reply_markup=keyb,
+        )
+    if type == "voice":
+        await message.reply_voice(
+            voice=file_id,
+            caption=data,
+            reply_markup=keyb,
+        )
 
 @app2.on_message(
     filters.command("delete", prefixes=USERBOT_PREFIX)

+ 7 - 6
wbb/utils/dbfunctions.py

@@ -451,7 +451,6 @@ async def save_captcha_solved(chat_id: int, user_id: int):
         upsert=True,
     )
 
-
 async def is_antiservice_on(chat_id: int) -> bool:
     chat = await antiservicedb.find_one({"chat_id": chat_id})
     if not chat:
@@ -498,15 +497,17 @@ async def get_welcome(chat_id: int) -> (str, str, str):
     data = await welcomedb.find_one({"chat_id": chat_id})
     if not data:
         return "", "", ""
+        
+    welcome = data.get("welcome", "")    
     raw_text = data.get("raw_text", "")
-    animation_id = data.get("animation_id", "")
-    photo_id = data.get("photo_id", "")
+    file_id = data.get("file_id", "")
+
+    return welcome, raw_text, file_id
 
-    return raw_text, animation_id, photo_id
 
+async def set_welcome(chat_id: int, welcome: str, raw_text: str, file_id: str):
+    update_data = {"welcome": welcome, "raw_text": raw_text, "file_id": file_id}
 
-async def set_welcome(chat_id: int, raw_text: str, animation_id: str, photo_id: str):
-    update_data = {"raw_text": raw_text, "animation_id": animation_id, "photo_id": photo_id}
     return await welcomedb.update_one(
         {"chat_id": chat_id},
         {"$set": update_data},