This commit is contained in:
REDNBLACK 2016-12-02 00:01:04 +03:00
parent af916a0f09
commit 03c8ea48c8
9 changed files with 50 additions and 74 deletions

View File

@ -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':

View File

@ -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

View File

@ -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)

View File

@ -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() \

View File

@ -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.

View File

@ -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)))

View File

@ -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:

View File

@ -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')))

View File

@ -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()