diff --git a/src/bot.py b/src/bot.py index 26a42cd..27e04a8 100644 --- a/src/bot.py +++ b/src/bot.py @@ -11,6 +11,7 @@ from src.service.chat_purge_queue import ChatPurgeQueue from src.service.data_learner import DataLearner from src.service.reply_generator import ReplyGenerator from src.service.links_checker import LinksChecker +from src.service.chance_manager import ChanceManager class Bot: @@ -22,11 +23,13 @@ class Bot: logging.info("Bot started") redis = Redis(config) + chance_manager = ChanceManager(redis) self.dispatcher.add_handler(MessageHandler(data_learner=DataLearner(), reply_generator=ReplyGenerator(), - links_checker=LinksChecker(redis))) - self.dispatcher.add_handler(CommandHandler()) + links_checker=LinksChecker(redis), + chance_manager=chance_manager)) + self.dispatcher.add_handler(CommandHandler(chance_manager=chance_manager)) self.dispatcher.add_handler(StatusHandler(chat_purge_queue=ChatPurgeQueue(self.updater.job_queue, redis))) if config['updates']['mode'] == 'polling': diff --git a/src/domain/abstract_entity.py b/src/domain/abstract_entity.py index 0493df4..c56c949 100644 --- a/src/domain/abstract_entity.py +++ b/src/domain/abstract_entity.py @@ -2,6 +2,7 @@ from abc import ABC class AbstractEntity(ABC): - def __init__(self, chat, message): - self.chat = chat + def __init__(self, message): + self.chat_id = message.chat.id + self.chat_type = message.chat.type self.message = message diff --git a/src/domain/command.py b/src/domain/command.py index a4f0be9..cba9542 100644 --- a/src/domain/command.py +++ b/src/domain/command.py @@ -2,8 +2,8 @@ from .abstract_entity import AbstractEntity class Command(AbstractEntity): - def __init__(self, chat, message): - super(Command, self).__init__(chat=chat, message=message) + def __init__(self, message): + super(Command, self).__init__(message) self.name = Command.parse_name(message) self.args = Command.parse_args(message) diff --git a/src/domain/message.py b/src/domain/message.py index 5d09c28..da74db0 100644 --- a/src/domain/message.py +++ b/src/domain/message.py @@ -7,8 +7,10 @@ from src.config import config class Message(AbstractEntity): - def __init__(self, chat, message): - super(Message, self).__init__(chat=chat, message=message) + def __init__(self, chance, message): + super(Message, self).__init__(message) + + self.chance = chance if self.has_text(): self.text = message.text @@ -63,7 +65,7 @@ class Message(AbstractEntity): 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', config['bot']['default_chance']) + return random.randint(0, 100) < self.chance def should_answer(self): return self.has_anchors() \ diff --git a/src/domain/status.py b/src/domain/status.py index a7a6f35..ea3d77c 100644 --- a/src/domain/status.py +++ b/src/domain/status.py @@ -4,8 +4,8 @@ from src.config import config class Status(AbstractEntity): - def __init__(self, chat, message): - super(Status, self).__init__(chat=chat, message=message) + def __init__(self, message): + super(Status, self).__init__(message) def is_bot_kicked(self): """Returns True if the bot was kicked from group. diff --git a/src/entity/chat.py b/src/entity/chat.py deleted file mode 100644 index 0fd05b2..0000000 --- a/src/entity/chat.py +++ /dev/null @@ -1,31 +0,0 @@ -import logging - -from orator.orm import Model -from orator.orm import has_many - -import src.entity.pair - - -class Chat(Model): - __fillable__ = ['telegram_id', 'chat_type', 'random_chance'] - - @has_many - def pairs(self): - return src.entity.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() - - @staticmethod - def get_chat(message): - return Chat.first_or_create(telegram_id=message.chat.id, chat_type=message.chat.type) - -# Events -Chat.created(lambda chat: logging.info("[Chat %s %s] Created with internal ID #%s" % - (chat.chat_type, chat.telegram_id, chat.id))) - -#TODO Fix -#Chat.updating(lambda chat: logging.info("[Chat %s %s New gab level is set to %s" % -# (chat.chat_type, chat.telegram_id, chat.random_chance))) \ No newline at end of file diff --git a/src/handler/command_handler.py b/src/handler/command_handler.py index 8525383..6596111 100644 --- a/src/handler/command_handler.py +++ b/src/handler/command_handler.py @@ -5,12 +5,12 @@ from telegram import Update from telegram.ext import Handler from src.domain.command import Command -from src.entity.chat import Chat +from src.entity.pair import Pair from .moderate_command import ModerateCommand class CommandHandler(Handler): - def __init__(self): + def __init__(self, chance_manager): super(CommandHandler, self).__init__(self.handle) self.commands = { 'start': self.__start_command, @@ -21,6 +21,7 @@ class CommandHandler(Handler): 'get_stats': self.__get_stats_command, 'moderate': self.__moderate_command } + self.chance_manager = chance_manager def check_update(self, update): if isinstance(update, Update) and update.message: @@ -36,8 +37,7 @@ class CommandHandler(Handler): return self.callback(dispatcher.bot, update, **optional_args) def handle(self, bot, update): - chat = Chat.get_chat(update.message) - command = Command(chat=chat, message=update.message) + command = Command(update.message) try: self.commands[command.name](bot, command) @@ -46,12 +46,12 @@ class CommandHandler(Handler): def __reply(self, bot, command, message): logging.debug("[Chat %s %s command] %s: %s" % - (command.chat.chat_type, - command.chat.telegram_id, + (command.chat_type, + command.chat_id, command.name, message)) - bot.send_message(chat_id=command.chat.telegram_id, + bot.send_message(chat_id=command.chat_id, reply_to_message_id=command.message.message_id, text=message) @@ -89,22 +89,24 @@ In 12 hours, I'll forget everything that have been learned in your chat, so you def __set_chance_command(self, bot, command): try: - random_chance = int(command.args[0]) + chance = int(command.args[0]) - if random_chance < 1 or random_chance > 50: + if chance < 1 or chance > 50: raise ValueError - command.chat.update(random_chance=random_chance) + self.chance_manager.set_chance(chat_id=command.chat_id, chance=chance) - self.__reply(bot, command, 'Set chance to: {}'.format(random_chance)) + self.__reply(bot, command, 'Set chance to: {}'.format(chance)) except (IndexError, ValueError): self.__reply(bot, command, 'Usage: /set_chance 1-50.') def __get_chance_command(self, bot, command): - self.__reply(bot, command, 'Current chance: {}'.format(command.chat.random_chance)) + self.__reply(bot, command, 'Current chance: {}' + .format(self.chance_manager.get_chance(command.chat_id))) def __get_stats_command(self, bot, command): - self.__reply(bot, command, 'Pairs: {}'.format(command.chat.pairs().count())) + self.__reply(bot, command, 'Pairs: {}' + .format(Pair.where('chat_id', command.chat_id).count())) def __moderate_command(self, bot, command): try: diff --git a/src/handler/message_handler.py b/src/handler/message_handler.py index c4dff25..f0a2f0a 100644 --- a/src/handler/message_handler.py +++ b/src/handler/message_handler.py @@ -4,13 +4,11 @@ from random import choice from src.config import config from telegram.ext import MessageHandler as ParentHandler, Filters from telegram import ChatAction - from src.domain.message import Message -from src.entity.chat import Chat class MessageHandler(ParentHandler): - def __init__(self, data_learner, reply_generator, links_checker): + def __init__(self, data_learner, reply_generator, links_checker, chance_manager): super(MessageHandler, self).__init__( Filters.text | Filters.sticker, self.handle) @@ -18,15 +16,16 @@ class MessageHandler(ParentHandler): self.data_learner = data_learner self.reply_generator = reply_generator self.links_checker = links_checker + self.chance_manager = chance_manager def handle(self, bot, update): - chat = Chat.get_chat(update.message) - message = Message(chat=chat, message=update.message) + chance = self.chance_manager.get_chance(update.message.chat.id) + message = Message(chance=chance, message=update.message) if message.has_text(): logging.debug("[Chat %s %s bare_text] %s" % - (message.chat.chat_type, - message.chat.telegram_id, + (message.chat_type, + message.chat_id, message.text)) if message.has_text() and not message.is_editing(): @@ -38,12 +37,12 @@ class MessageHandler(ParentHandler): should_answer = message.should_answer() if should_answer: - bot.send_chat_action(chat_id=message.chat.telegram_id, action=ChatAction.TYPING) + bot.send_chat_action(chat_id=message.chat_id, action=ChatAction.TYPING) self.data_learner.learn(message) - if message.has_links() and self.links_checker.check(message.chat.telegram_id, message.links): - bot.send_sticker(chat_id=message.chat.telegram_id, + if message.has_links() and self.links_checker.check(message.chat_id, message.links): + bot.send_sticker(chat_id=message.chat_id, reply_to_message_id=message.message.message_id, sticker=choice(config.getlist('links', 'stickers'))) @@ -52,20 +51,20 @@ class MessageHandler(ParentHandler): reply_id = None if not message.is_reply_to_bot() else message.message.message_id logging.debug("[Chat %s %s answer] %s" % - (message.chat.chat_type, - message.chat.telegram_id, + (message.chat_type, + message.chat_id, text)) - bot.send_message(chat_id=message.chat.telegram_id, + bot.send_message(chat_id=message.chat_id, reply_to_message_id=reply_id, text=text) def __process_sticker(self, bot, message): if message.should_answer(): logging.debug("[Chat %s %s spam_sticker]" % - (message.chat.chat_type, - message.chat.telegram_id)) + (message.chat_type, + message.chat_id)) - bot.send_sticker(chat_id=message.chat.telegram_id, + bot.send_sticker(chat_id=message.chat_id, reply_to_message_id=message.message.message_id, sticker=choice(config.getlist('bot', 'spam_stickers'))) diff --git a/src/handler/moderate_command.py b/src/handler/moderate_command.py index 9752898..b6bdeb3 100644 --- a/src/handler/moderate_command.py +++ b/src/handler/moderate_command.py @@ -33,7 +33,7 @@ class ModerateCommand: return [] to_keep = [] - for pair in self.__find_pairs(self.entity.chat.id, list(found_words.keys())): + for pair in self.__find_pairs(self.entity.chat_id, list(found_words.keys())): if pair.first_id in found_words: to_keep.append(pair.first_id) if pair.second_id in found_words: @@ -47,13 +47,13 @@ class ModerateCommand: user_id = self.entity.message.from_user.id admin_ids = list(map( lambda member: member.user.id, - self.bot.get_chat_administrators(chat_id=self.entity.chat.telegram_id) + self.bot.get_chat_administrators(chat_id=self.entity.chat_id) )) return user_id in admin_ids def remove_word(self, word_id): - pairs_ids = self.__find_pairs(self.entity.chat.id, [word_id]).lists('id') + pairs_ids = self.__find_pairs(self.entity.chat_id, [word_id]).lists('id') Pair.where_in('id', pairs_ids).delete() Reply.where_in('pair_id', pairs_ids).delete()