Merge branch 'master' into redis

This commit is contained in:
REDNBLACK 2016-12-28 21:30:01 +03:00
commit ef399285df
11 changed files with 85 additions and 18 deletions

View File

@ -7,6 +7,9 @@ from src.handler import *
class Bot: class Bot:
"""
Main initializer and dispatcher of messages
"""
def __init__(self): def __init__(self):
self.updater = Updater(token=config['bot']['token']) self.updater = Updater(token=config['bot']['token'])
self.dispatcher = self.updater.dispatcher self.dispatcher = self.updater.dispatcher

View File

@ -2,13 +2,17 @@ from abc import ABC
class AbstractEntity(ABC): class AbstractEntity(ABC):
"""
Base class for all message entities
"""
def __init__(self, message): def __init__(self, message):
self.chat_id = message.chat.id self.chat_id = message.chat.id
self.chat_type = message.chat.type self.chat_type = message.chat.type
self.message = message self.message = message
def is_private(self): def is_private(self):
"""Returns True if the message is private. """Returns True if chat type is private.
""" """
return self.message.chat.type == 'private' return self.message.chat.type == 'private'

View File

@ -2,6 +2,9 @@ from .abstract_entity import AbstractEntity
class Command(AbstractEntity): class Command(AbstractEntity):
"""
Special class for message which contains command
"""
def __init__(self, message): def __init__(self, message):
super(Command, self).__init__(message) super(Command, self).__init__(message)
self.name = Command.parse_name(message) self.name = Command.parse_name(message)
@ -9,8 +12,18 @@ class Command(AbstractEntity):
@staticmethod @staticmethod
def parse_name(message): def parse_name(message):
"""
Parses command name from given message
:param message: Telegram message object
:return: Name of command
"""
return message.text[1:].split(' ')[0].split('@')[0] return message.text[1:].split(' ')[0].split('@')[0]
@staticmethod @staticmethod
def parse_args(message): def parse_args(message):
"""
Parses command args from given message
:param message: Telegram message object
:return: List of command args
"""
return message.text.split()[1:] return message.text.split()[1:]

View File

@ -1,16 +1,19 @@
import random import random
import re
from .abstract_entity import AbstractEntity from .abstract_entity import AbstractEntity
from src.utils import deep_get_attr from src.utils import deep_get_attr
from src.config import config from src.config import config
class Message(AbstractEntity): class Message(AbstractEntity):
"""
Basic message entity
"""
def __init__(self, chance, message): def __init__(self, chance, message):
super(Message, self).__init__(message) super(Message, self).__init__(message)
self.chance = chance self.chance = chance
self.entities = message.entities self.entities = message.entities
self.anchors = config.getlist('bot', 'anchors')
if self.has_text(): if self.has_text():
self.text = message.text self.text = message.text
@ -18,39 +21,48 @@ class Message(AbstractEntity):
self.text = '' self.text = ''
def has_text(self): def has_text(self):
"""Returns True if the message has text. """
Returns True if the message has text.
""" """
return self.message.text.strip() != '' return self.message.text.strip() != ''
def is_sticker(self): def is_sticker(self):
"""Returns True if the message is a sticker. """
Returns True if the message is a sticker.
""" """
return self.message.sticker is not None return self.message.sticker is not None
def has_entities(self): def has_entities(self):
"""Returns True if the message has entities (attachments). """
Returns True if the message has entities (attachments).
""" """
return self.entities is not None return self.entities is not None
def has_anchors(self): def has_anchors(self):
"""Returns True if the message contains at least one anchor from anchors config.
""" """
anchors = config.getlist('bot', 'anchors') Returns True if the message contains at least one anchor from anchors config.
return self.has_text() and any(a in self.message.text.split(' ') for a in anchors) """
return self.has_text() and any(a in self.message.text.split(' ') for a in self.anchors)
def is_reply_to_bot(self): def is_reply_to_bot(self):
"""Returns True if the message is a reply to bot. """
Returns True if the message is a reply to bot.
""" """
user_name = deep_get_attr(self.message, 'reply_to_message.from_user.username') user_name = deep_get_attr(self.message, 'reply_to_message.from_user.username')
return user_name == config['bot']['name'] return user_name == config['bot']['name']
def is_random_answer(self): def is_random_answer(self):
"""Returns True if reply chance for this chat is high enough """
Returns True if reply chance for this chat is high enough
""" """
return random.randint(0, 100) < self.chance return random.randint(0, 100) < self.chance
def should_answer(self): def should_answer(self):
"""
Returns True if bot should answer to this message
:return: Should answer or not
"""
return self.has_anchors() \ return self.has_anchors() \
or self.is_private() \ or self.is_private() \
or self.is_reply_to_bot() \ or self.is_reply_to_bot() \

View File

@ -4,20 +4,24 @@ from src.config import config
class Status(AbstractEntity): class Status(AbstractEntity):
bot_name = config['bot']['name'] """
Special class for message which contains info about status change
"""
def __init__(self, message): def __init__(self, message):
super(Status, self).__init__(message) super(Status, self).__init__(message)
self.bot_name = config['bot']['name']
def is_bot_kicked(self): def is_bot_kicked(self):
"""Returns True if the bot was kicked from group. """
Returns True if the bot was kicked from group.
""" """
user_name = deep_get_attr(self.message, 'left_chat_member.username') user_name = deep_get_attr(self.message, 'left_chat_member.username')
return user_name == self.bot_name return user_name == self.bot_name
def is_bot_added(self): def is_bot_added(self):
"""Returns True if the bot was added to group. """
Returns True if the bot was added to group.
""" """
user_name = deep_get_attr(self.message, 'new_chat_member.username') user_name = deep_get_attr(self.message, 'new_chat_member.username')

View File

@ -1,7 +1,7 @@
import logging import logging
from random import choice from random import choice
from src.config import config, data_learner, reply_generator, media_checker, chance_manager from src.config import config, data_learner, reply_generator, media_checker, chance_repository
from telegram.ext import MessageHandler as ParentHandler, Filters from telegram.ext import MessageHandler as ParentHandler, Filters
from telegram import ChatAction from telegram import ChatAction
from src.domain.message import Message from src.domain.message import Message
@ -15,12 +15,12 @@ class MessageHandler(ParentHandler):
self.data_learner = data_learner self.data_learner = data_learner
self.reply_generator = reply_generator self.reply_generator = reply_generator
self.media_checker = media_checker self.media_checker = media_checker
self.chance_manager = chance_manager self.chance_repository = chance_repository
self.spam_stickers = config.getlist('bot', 'spam_stickers') self.spam_stickers = config.getlist('bot', 'spam_stickers')
self.media_checker_messages = config.getlist('media_checker', 'messages') self.media_checker_messages = config.getlist('media_checker', 'messages')
def handle(self, bot, update): def handle(self, bot, update):
chance = self.chance_manager.get_chance(update.message.chat.id) chance = self.chance_repository.get(update.message.chat.id)
message = Message(chance=chance, message=update.message) message = Message(chance=chance, message=update.message)
self.__check_media_uniqueness(bot, message) self.__check_media_uniqueness(bot, message)

View File

@ -2,6 +2,9 @@ import redis
class Redis: class Redis:
"""
Small redis wrapper, to simplify work with connection pool
"""
def __init__(self, config): def __init__(self, config):
self.pool = redis.ConnectionPool(host=config['redis']['host'], self.pool = redis.ConnectionPool(host=config['redis']['host'],
port=config.getint('redis', 'port'), port=config.getint('redis', 'port'),

View File

@ -6,6 +6,9 @@ from src.config import config, trigram_repository, job_repository
class ChatPurgeQueue: class ChatPurgeQueue:
"""
Scheduling and execution of chat purge
"""
def __init__(self): def __init__(self):
self.queue = None self.queue = None
self.jobs = {} self.jobs = {}
@ -21,6 +24,12 @@ class ChatPurgeQueue:
return self return self
def add(self, chat_id, interval=None, db=True): def add(self, chat_id, interval=None, db=True):
"""
Schedules purge of chat data
:param chat_id: ID of chat
:param interval: Interval in seconds
:param db: Should be added to db or not
"""
interval = interval if interval is not None else self.default_interval interval = interval if interval is not None else self.default_interval
scheduled_at = datetime.now() + timedelta(seconds=interval) scheduled_at = datetime.now() + timedelta(seconds=interval)
@ -35,6 +44,10 @@ class ChatPurgeQueue:
self.job_repository.add(chat_id, scheduled_at) self.job_repository.add(chat_id, scheduled_at)
def remove(self, chat_id): def remove(self, chat_id):
"""
Removes scheduled purge job from queue
:param chat_id: ID of chat
"""
if chat_id not in self.jobs: if chat_id not in self.jobs:
return return

View File

@ -7,6 +7,10 @@ class DataLearner:
self.tokenizer = tokenizer self.tokenizer = tokenizer
def learn(self, message): def learn(self, message):
"""
Split message to trigrams and write to DB
:param message: Message
"""
words = self.tokenizer.extract_words(message) words = self.tokenizer.extract_words(message)
trigrams = self.tokenizer.split_to_trigrams(words) trigrams = self.tokenizer.split_to_trigrams(words)

View File

@ -4,6 +4,9 @@ from urllib.parse import urlparse
class MediaUniquenessChecker: class MediaUniquenessChecker:
"""
Checks message links and photos for uniqueness
"""
def __init__(self): def __init__(self):
self.media_repository = media_repository self.media_repository = media_repository
@ -19,7 +22,7 @@ class MediaUniquenessChecker:
def __extract_media(self, message): def __extract_media(self, message):
media = [] media = []
for entity in filter(lambda e: e.type == 'url', message.message.entities): for entity in filter(lambda e: e.type == 'url', message.entities):
link = self.__prettify(message.text[entity.offset:entity.length + entity.offset]) link = self.__prettify(message.text[entity.offset:entity.length + entity.offset])
media.append(link) media.append(link)

View File

@ -3,6 +3,9 @@ from src.utils import strings_has_equal_letters, capitalize
class ReplyGenerator: class ReplyGenerator:
"""
Handles generation of responses for user message
"""
def __init__(self): def __init__(self):
self.redis = redis self.redis = redis
self.tokenizer = tokenizer self.tokenizer = tokenizer
@ -16,6 +19,11 @@ class ReplyGenerator:
self.end_sentence = config['grammar']['end_sentence'] self.end_sentence = config['grammar']['end_sentence']
def generate(self, message): def generate(self, message):
"""
Generates response based on message words
:param message: Message
:return: Response or empty string, if generated response equals to user message
"""
words = self.tokenizer.extract_words(message) words = self.tokenizer.extract_words(message)
pairs = [trigram[:-1] for trigram in self.tokenizer.split_to_trigrams(words)] pairs = [trigram[:-1] for trigram in self.tokenizer.split_to_trigrams(words)]
messages = [self.__generate_best_message(chat_id=message.chat_id, pair=pair) for pair in pairs] messages = [self.__generate_best_message(chat_id=message.chat_id, pair=pair) for pair in pairs]