Refactor note saving, fix a bug causing no note to be detected

PaulSonOfLars · Jun 18, 2018 · 08b0a41 · 08b0a41
1 parent 20ea344
commit 08b0a41
Show file tree

Hide file tree

Showing 2 changed files with 50 additions and 72 deletions.
diff --git a/tg_bot/modules/helper_funcs/msg_types.py b/tg_bot/modules/helper_funcs/msg_types.py
@@ -17,7 +17,7 @@ class Types(IntEnum):
     VIDEO = 7
 
 
-def get_note_type(msg: Message, replied=False):
+def get_note_type(msg: Message):
     data_type = None
     content = None
     text = ""
@@ -36,47 +36,45 @@ def get_note_type(msg: Message, replied=False):
         else:
             data_type = Types.TEXT
 
-    elif len(args) >= 2 and replied and msg.reply_to_message and msg.reply_to_message.text:  # not caption, text
-        text, buttons = button_markdown_parser(msg.reply_to_message.text,
-                                               entities=msg.reply_to_message.parse_entities())
-        if buttons:
-            data_type = Types.BUTTON_TEXT
-        else:
-            data_type = Types.TEXT
-
-    elif msg.reply_to_message and msg.reply_to_message.sticker:
-        content = msg.reply_to_message.sticker.file_id
-        data_type = Types.STICKER
-
-    elif msg.reply_to_message and msg.reply_to_message.document:
-        content = msg.reply_to_message.document.file_id
-        text, buttons = button_markdown_parser(msg.reply_to_message.caption,
-                                               entities=msg.reply_to_message.parse_entities())
-        data_type = Types.DOCUMENT
-
-    elif msg.reply_to_message and msg.reply_to_message.photo:
-        content = msg.reply_to_message.photo[-1].file_id  # last elem = best quality
-        text, buttons = button_markdown_parser(msg.reply_to_message.caption,
-                                               entities=msg.reply_to_message.parse_entities())
-        data_type = Types.PHOTO
-
-    elif msg.reply_to_message and msg.reply_to_message.audio:
-        content = msg.reply_to_message.audio.file_id
-        text, buttons = button_markdown_parser(msg.reply_to_message.caption,
-                                               entities=msg.reply_to_message.parse_entities())
-        data_type = Types.AUDIO
-
-    elif msg.reply_to_message and msg.reply_to_message.voice:
-        content = msg.reply_to_message.voice.file_id
-        text, buttons = button_markdown_parser(msg.reply_to_message.caption,
-                                               entities=msg.reply_to_message.parse_entities())
-        data_type = Types.VOICE
-
-    elif msg.reply_to_message and msg.reply_to_message.video:
-        content = msg.reply_to_message.video.file_id
-        text, buttons = button_markdown_parser(msg.reply_to_message.caption,
-                                               entities=msg.reply_to_message.parse_entities())
-        data_type = Types.VIDEO
+    elif msg.reply_to_message:
+        entities = msg.reply_to_message.parse_entities()
+        msgtext = msg.reply_to_message.text or msg.reply_to_message.caption
+        if len(args) >= 2 and msg.reply_to_message.text:  # not caption, text
+            text, buttons = button_markdown_parser(msgtext,
+                                                   entities=entities)
+            if buttons:
+                data_type = Types.BUTTON_TEXT
+            else:
+                data_type = Types.TEXT
+
+        elif msg.reply_to_message.sticker:
+            content = msg.reply_to_message.sticker.file_id
+            data_type = Types.STICKER
+
+        elif msg.reply_to_message.document:
+            content = msg.reply_to_message.document.file_id
+            text, buttons = button_markdown_parser(msgtext, entities=entities)
+            data_type = Types.DOCUMENT
+
+        elif msg.reply_to_message.photo:
+            content = msg.reply_to_message.photo[-1].file_id  # last elem = best quality
+            text, buttons = button_markdown_parser(msgtext, entities=entities)
+            data_type = Types.PHOTO
+
+        elif msg.reply_to_message.audio:
+            content = msg.reply_to_message.audio.file_id
+            text, buttons = button_markdown_parser(msgtext, entities=entities)
+            data_type = Types.AUDIO
+
+        elif msg.reply_to_message.voice:
+            content = msg.reply_to_message.voice.file_id
+            text, buttons = button_markdown_parser(msgtext, entities=entities)
+            data_type = Types.VOICE
+
+        elif msg.reply_to_message.video:
+            content = msg.reply_to_message.video.file_id
+            text, buttons = button_markdown_parser(msgtext, entities=entities)
+            data_type = Types.VIDEO
 
     return note_name, text, data_type, content, buttons
 

diff --git a/tg_bot/modules/notes.py b/tg_bot/modules/notes.py
@@ -5,7 +5,7 @@
 from telegram import MAX_MESSAGE_LENGTH, ParseMode, InlineKeyboardMarkup
 from telegram import Message, Update, Bot
 from telegram.error import BadRequest
-from telegram.ext import CommandHandler, RegexHandler, Filters
+from telegram.ext import CommandHandler, RegexHandler
 from telegram.ext.dispatcher import run_async
 from telegram.utils.helpers import escape_markdown
 
@@ -127,23 +127,24 @@ def hash_get(bot: Bot, update: Update):
     get(bot, update, no_hash, show_none=False)
 
 
-# TODO: FIX THIS
 @run_async
 @user_admin
-def save_replied(bot: Bot, update: Update):
+def save(bot: Bot, update: Update):
     chat_id = update.effective_chat.id
-    msg = update.effective_message
+    msg = update.effective_message  # type: Optional[Message]
 
-    notename, text, data_type, content, buttons = get_note_type(msg, replied=False)
+    note_name, text, data_type, content, buttons = get_note_type(msg)
 
     if data_type is None:
         msg.reply_text("Dude, there's no note")
         return
 
-    sql.add_note_to_db(chat_id, notename, text, data_type, buttons, content)
-    msg.reply_text("Yas! Added replied message {}".format(notename))
+    sql.add_note_to_db(chat_id, note_name, text, data_type, buttons=buttons, file=content)
+
+    msg.reply_text(
+        "Yas! Added {note_name}.\nGet it with /get {note_name}, or #{note_name}".format(note_name=note_name))
 
-    if msg.reply_to_message.from_user.is_bot:
+    if msg.reply_to_message and msg.reply_to_message.from_user.is_bot:
         if text:
             msg.reply_text("Seems like you're trying to save a message from a bot. Unfortunately, "
                            "bots can't forward bot messages, so I can't save the exact message. "
@@ -157,25 +158,6 @@ def save_replied(bot: Bot, update: Update):
         return
 
 
-@run_async
-@user_admin
-def save(bot: Bot, update: Update):
-    chat_id = update.effective_chat.id
-    msg = update.effective_message  # type: Optional[Message]
-    raw_text = msg.text
-
-    note_name, text, data_type, content, buttons = get_note_type(msg)
-
-    if data_type is None:
-        msg.reply_text("Dude, there's no note")
-        return
-
-    sql.add_note_to_db(chat_id, note_name, text, data_type, buttons=buttons, file=content)
-
-    msg.reply_text(
-        "Yas! Added {note_name}.\nGet it with /get {note_name}, or #{note_name}".format(note_name=note_name))
-
-
 @run_async
 @user_admin
 def clear(bot: Bot, update: Update, args: List[str]):
@@ -265,15 +247,13 @@ def __chat_settings__(chat_id, user_id):
 GET_HANDLER = CommandHandler("get", cmd_get, pass_args=True)
 HASH_GET_HANDLER = RegexHandler(r"^#[^\s]+", hash_get)
 
-SAVE_HANDLER = CommandHandler("save", save, filters=~Filters.reply)
-REPL_SAVE_HANDLER = CommandHandler("save", save_replied, filters=Filters.reply)
+SAVE_HANDLER = CommandHandler("save", save)
 DELETE_HANDLER = CommandHandler("clear", clear, pass_args=True)
 
 LIST_HANDLER = DisableAbleCommandHandler(["notes", "saved"], list_notes, admin_ok=True)
 
 dispatcher.add_handler(GET_HANDLER)
 dispatcher.add_handler(SAVE_HANDLER)
-dispatcher.add_handler(REPL_SAVE_HANDLER)
 dispatcher.add_handler(LIST_HANDLER)
 dispatcher.add_handler(DELETE_HANDLER)
 dispatcher.add_handler(HASH_GET_HANDLER)