Implemented #6 and insignificant refactoring

src/bot.py

@@ -1,10 +1,11 @@
 import logging
 
 from telegram.ext import Updater
-from telegram.ext import MessageHandler, Filters
+from telegram.ext import MessageHandler, CommandHandler, Filters
 from orator.orm import Model
 from orator import DatabaseManager
 from src.message import Message
+from src.command_manager import CommandManager
 
 
 class Bot:
@@ -12,14 +13,27 @@ class Bot:
         self.config = config
         self.updater = Updater(token=config['bot']['token'])
         self.dispatcher = self.updater.dispatcher
+        self.command_manager = CommandManager()
 
         Model.set_connection_resolver(DatabaseManager({'db': config['db']}))
 
-    def handler(self, bot, update):
+    def message_handler(self, bot, update):
         Message(bot=bot, message=update.message, config=self.config).process()
 
+    def command_handler(self, bot, update, args):
+        try:
+            command = update.message.text.strip('/').split(' ')[0]
+
+            self.command_manager.handle(bot=bot, update=update, command=command, args=args)
+        except (IndexError, ValueError):
+            update.message.reply_text('Invalid command!')
+
     def run(self):
         logging.info("Bot started")
-        handler = MessageHandler([Filters.text], self.handler)
-        self.dispatcher.add_handler(handler)
+        message_handler = MessageHandler(Filters.text, self.message_handler)
+        command_handler = CommandHandler('set_chance', self.command_handler, pass_args=True)
+
+        self.dispatcher.add_handler(message_handler)
+        self.dispatcher.add_handler(command_handler)
         self.updater.start_polling()
+        self.updater.idle()

src/command_manager.py

@@ -0,0 +1,30 @@
+from src.domain.chat import Chat
+
+
+class CommandManager:
+    def __init__(self):
+        self.commands = {
+            'set_chance': self.__set_chance_command
+        }
+
+    def handle(self, bot, update, command, args):
+        if command in self.commands:
+            method = self.commands[command]
+            method(update, args)
+        else:
+            update.message.reply_text('Invalid command!')
+
+    def __set_chance_command(self, update, args):
+        try:
+            random_chance = int(args[0])
+
+            if random_chance < 1 or random_chance > 50:
+                raise ValueError
+
+            chat = Chat.get_chat(update.message)
+            chat.random_chance = random_chance
+            chat.save()
+
+            update.message.reply_text('Set chance to %d' % random_chance)
+        except (IndexError, ValueError):
+            update.message.reply_text('Usage: /set_chance 1-50')

src/domain/chat.py

@@ -12,10 +12,10 @@ class Chat(Model):
     def pairs(self):
         return src.domain.pair.Pair
 
-    def migrate_to_chat_id(self, new_id):
-        logging.info("[Chat %s %s] Migrating ID to %s" % (self.chat_type, self.telegram_id, new_id))
-        self.telegram_id = new_id
-        self.save()
+    # def migrate_to_chat_id(self, new_id):
+    #     logging.info("[Chat %s %s] Migrating ID to %s" % (self.chat_type, self.telegram_id, new_id))
+    #     self.telegram_id = new_id
+    #     self.save()
 
     @staticmethod
     def get_chat(message):

src/domain/pair.py

@@ -62,29 +62,30 @@ class Pair(Model):
 
         while any(word for word in words):
             trigram = words[:3]
-            words.pop(0)
-            trigram_word_ids = list(map(
+            first_word_id, second_word_id, *third_word_id = list(map(
                 lambda x: None if x is None else src.domain.word.Word.where('word', x).first().id,
                 trigram
             ))
+            third_word_id = None if len(third_word_id) == 0 else third_word_id[0]
+
+            words.pop(0)
 
             pair = Pair.where('chat_id', message.chat.id)\
-                .where('first_id', trigram_word_ids[0])\
-                .where('second_id', trigram_word_ids[1])\
+                .where('first_id', first_word_id)\
+                .where('second_id', second_word_id)\
                 .first()
             if pair is None:
                 pair = Pair.create(chat_id=message.chat.id,
-                                   first_id=trigram_word_ids[0],
-                                   second_id=trigram_word_ids[1])
+                                   first_id=first_word_id,
+                                   second_id=second_word_id)
 
-            last_trigram_id = trigram_word_ids[2] if len(trigram_word_ids) == 3 else None
-            reply = pair.replies().where('word_id', last_trigram_id).first()
+            reply = pair.replies().where('word_id', third_word_id).first()
 
             if reply is not None:
                 reply.count += 1
                 reply.save()
             else:
-                src.domain.reply.Reply.create(pair_id=pair.id, word_id=last_trigram_id)
+                src.domain.reply.Reply.create(pair_id=pair.id, word_id=third_word_id)
 
     @staticmethod
     def __generate_sentence(message, word_ids):

src/message.py

@@ -3,32 +3,27 @@ import random
 
 import src.domain.chat as chat
 import src.domain.pair as pair
+from src.utils import deep_get_attr
 
 
 class Message:
     def __init__(self, bot, message, config):
-        self.bot                     = bot
-        self.message                 = message
-        self.config                  = config
-        self.chat                    = chat.Chat.get_chat(message)
-        self.chat.migrate_to_chat_id = message.migrate_to_chat_id
+        self.bot     = bot
+        self.message = message
+        self.config  = config
+        self.chat    = chat.Chat.get_chat(message)
 
         if self.has_text():
             logging.debug("[chat %s %s bare_text] %s" %
                           (self.chat.chat_type, 
                            self.chat.telegram_id, 
                            self.message.text))
-            self.text    = message.text
-            self.words   = self.__get_words()
-            self.command = self.__get_command() if self.text[0] == '/' else None
+            self.text  = message.text
+            self.words = self.__get_words()
 
     def process(self):
-        if self.has_text() and not self.is_editing():
-            # TODO Implement
-            if self.is_command():
-                pass
-            else:
-                return self.__process_message()
+        if self.has_text() and not (self.is_editing() or self.is_command()):
+            return self.__process_message()
 
     def has_text(self):
         """Returns True if the message has text.
@@ -46,6 +41,8 @@ class Message:
         return self.message.entities is not None
 
     def has_anchors(self):
+        """Returns True if the message contains at least one anchor from anchors config.
+        """
         anchors = self.config['bot']['anchors'].split(',')
         return self.has_text() and \
                (any(x in self.message.text.split(' ') for x in anchors))
@@ -58,22 +55,22 @@ class Message:
     def is_reply_to_bot(self):
         """Returns True if the message is a reply to bot.
         """
-        reply_to_message = getattr(self.message, 'reply_to_message', None)
-        from_user        = getattr(reply_to_message, 'from_user', None)
-        user_name        = getattr(from_user, 'username', None)
+        user_name = deep_get_attr(self.message, 'reply_to_message.from_user.username', None)
 
         return user_name == self.config['bot']['name']
 
     def is_random_answer(self):
+        """Returns True if reply chance for this chat is high enough
+        """
         return random.randint(0, 100) < getattr(self.chat, 'random_chance', 5)
 
     def is_command(self):
         """Returns True if the message is a command (`/start`, `/do_stuff`).
         """
-        return self.command is not None
+        return self.has_text() and self.text[0] == '/'
 
     def __answer(self, message):
-        logging.debug("[Chat %s %s answer] %s" %
+        logging.info("[Chat %s %s answer] %s" %
                       (self.chat.chat_type, 
                        self.chat.telegram_id, 
                        message))
@@ -81,7 +78,7 @@ class Message:
         self.bot.sendMessage(chat_id=self.chat.telegram_id, text=message)
 
     def __reply(self, message):
-        logging.debug("[Chat %s %s reply] %s" %
+        logging.info("[Chat %s %s reply] %s" %
                       (self.chat.chat_type, 
                        self.chat.telegram_id, 
                        message))
@@ -109,13 +106,9 @@ class Message:
         
         result = list(filter(None, map(lambda x: x.lower(), ''.join(text).split(' '))))
         
-        logging.debug("[chat %s %s get_words] %s" %
+        logging.debug("[Chat %s %s get_words] %s" %
                       (self.chat.chat_type,
                        self.chat.telegram_id, 
-                       result))
+                       ' '.join(result)))
 
         return result
-
-    # TODO Implement
-    def __get_command(self):
-        return None

src/utils.py

@@ -7,3 +7,15 @@ def capitalize(string):
 
 def random_element(xlist):
     return random.choice(xlist) if len(xlist) > 0 else None
+
+
+def deep_get_attr(obj, attr, default=None):
+    attributes = attr.split(".")
+    for i in attributes:
+        try:
+            obj = getattr(obj, i)
+            if callable(obj):
+                obj = obj()
+        except AttributeError:
+            return default
+    return obj