Пишем диалоговые Telegram-боты на Питоне

Думаю, всем здесь в той или иной мере известен мессенджер Telegram. Создатель заявляет, что это самый безопасный мессенджер с убойным алгоритмом шифрования собственной разработки, но нас, разработчиков, конечно же, куда сильнее интересует другое. Боты!

Тема эта, конечно, не раз поднималась на Хабре: ботов писали на Python с tornado, Node.js, Ruby со специальным гемом, Ruby on Rails, C#, C# с WCF и даже PHP; ботов писали для RSS-каналов, мониторинга сайтов, удалённого включения компьютера и, вероятно, для многого, многого другого.

И всё же я возьму на себя смелость изъездить эту тему ещё раз и вдобавок к этому показать немного магии Питона. Мы будем писать фреймворк™ для удобного написания нетривиальных диалоговых ботов на основе пакета python-telegram-bot.

Как зачать бота?

На этот вопрос лучше всего отвечает официальная документация. Выглядит процесс примерно так:



Просто, не правда ли? (Будьте благоразумны и не занимайте хорошие никнеймы без убедительной причины!)

Самый простой бот

Сперва глянем в туториал нашего базового пакета, чтобы понять, с чего начинается простенький бот. Следующий код

# -*- coding: utf-8 -*-
from telegram.ext import Updater         # пакет называется python-telegram-bot, но Python-
from telegram.ext import CommandHandler  # модуль почему-то просто telegram ¯\_(ツ)_/¯

def start(bot, update):
    # подробнее об объекте update: https://core.telegram.org/bots/api#update
    bot.sendMessage(chat_id=update.message.chat_id, text="Здравствуйте.")

updater = Updater(token='TOKEN')  # тут токен, который выдал вам Ботский Отец!

start_handler = CommandHandler('start', start)  # этот обработчик реагирует
                                                # только на команду /start

updater.dispatcher.add_handler(start_handler)   # регистрируем в госреестре обработчиков
updater.start_polling()  # поехали!

создаёт бота, который сухо отвечает «Здравствуйте.» при нажатии на кнопку Start (или ручном вводе команды /start) и многозначительно молчит при любых последующих действиях с вашей стороны.

Соответственно, если мы захотим повесить обработчики любых текстовых сообщений или любых команд, нам нужно будет написать

from telegram.ext import Filters, MessageHandler

def handle_text(bot, update):
    # ...

def handle_command(bot, update):
    # ...

# MessageHandler -- более универсальный обработчик, который берёт на вход фильтр
text_handler = MessageHandler(Filters.text, self.handle_text)
command_handler = MessageHandler(Filters.command, self.handle_command)
# регистрируем свеженькие обработчики в диспетчере
updater.dispatcher.add_handler(text_handler)     # без регистрации будет работать, 
updater.dispatcher.add_handler(command_handler)  # но не больше трёх месяцев (шутка)

(За дальнейшими подробностями с чистой совестью отсылаю к документации python-telegram-bot.)

Нагруженные этим теоретическим минимумом, мы можем наконец подумать, как нам писать своего нетривиального бота. Для начала давайте вернёмся к постановке задачи. Под диалоговым ботом я подразумеваю бота, который главным образом ведёт обычный текстовый диалог с пользователем — с вопросами, ответами, нелинейным сюжетом, разочаровывающими концовками и всем в таком духе (играли в «Бесконечное лето»?) Напротив, не попадают в сферу наших текущих интересов боты, разным образом расширяющие функционал Telegram (вроде бота для лайков); соответственно, мы опустим добавление всяких плюшек вроде инлайнового режима, игр, обновления элементов управления на лету и всего такого прочего.

Проблема сложных диалоговых ботов в том, что нетривиальный диалог требует хранения состояния. Работа асинхронных диалогов требует постоянных прерываний на ожидание сообщения от пользователя; состояние нужно сохранять, потом восстанавливать, прыгать к коду, ответственному за обработку очередного сообщения, и так далее; в общем, организация кода становится проблемой довольно угнетающей. Прервать, продолжить… ничего не напоминает? Что ж, посмотрим, как обозначенную проблему можно изящнейше обойти с помощью магии yield.

50 оттенков yield

Что мы знаем про yield? Ну, все мы знаем, что это такая штука, чтобы писать генераторы, то есть этакие ленивые и потенциально бесконечные списки:

def fibonacci():
    a, b = 1, 1
    while True:
        yield a
        a, b = b, a + b

f = fibonacci()

Теперь объект f — это такая волшебная коробка; стоит сунуть в неё руку написать next(f), и мы получим очередное число Фибоначчи, но стоит перевернуть её написать list(f), как мы уйдём в бесконечный цикл, который, скорее всего, закончится трагической смертью системы от нехватки оперативной памяти.

Мы знаем, что генераторы — это быстро, удобно и очень в стиле Python. У нас есть модуль itertools, предлагающий генераторы на любой вкус и цвет. Но у нас есть кое-что ещё.

Куда менее известными навыками слова yield являются способности… возвращать значения и бросать исключения! Да-да, если мы напишем:

f.throw(Exception)

То вычисление чисел Фибоначчи оборвётся самым трагическим образом — исключением в строчке с yield.

В свою очередь вызов f.send(something) заставит конструкцию yield вернуть значение, а потом сразу вернёт next(f). Достаточно приравнять yield переменную, чтобы переданное значение поймать:

def doubler(start):
    while True:
        start = yield (start, start)

d = doubler(42)
print(next(d))  # выводит (42, 42)
print(next(d))  # выводит (None, None), потому что мы забыли что-либо передать!
print(d.send(43))  # выводит (43, 43) -- yield вернул 43, и тут же сработал следующий yield

Но и это ещё не всё. Начиная с Python 3.3, генераторы умеют делегировать выполнение друг другу с помощью конструкции yield from: вместо

def concatenate(iterable1, iterable2):
    for item in iterable1:
        yield item
    for item in iterable2:
        yield item

она позволяет нам писать

def concatenate(iterable1, iterable2):
    yield from iterable1
    yield from iterable2

Но было бы преуменьшением сказать, что yield from позволяет нам лишь сэкономить строки кода на циклах. Дело в том, что она также заботится о send и throw — при вызове они будут взаимодействовать не с функцией concatenate, а с одним из двух генераторов, которым она передаёт управление. (Если это окажутся не генераторы… ну упс.)

А ещё yield from тоже умеет возвращать значение: для этого функциям-генераторам вернули право на нетривиальный (то есть возвращающий что-то, а не просто заканчивающий выполнение) return:

def despaired_person():
    yield None
    yield None
    yield None
    return "I'm tired of my uselessness"

def despair_expresser():
    result = yield from despaired_person()
    print(result)

print(list(f())) # печатает
# I'm tired of my uselessness
# [None, None, None]

К чему я всё это? Ах да. Эти фокусы, вместе взятые, позволят нам легко и естественно писать наших диалоговых ботов.

Пишем обёртку

Итак, пусть диалог с каждым пользователем ведётся генератором. yield будет выдавать наружу сообщение, которое надо отправить пользователю, и возвращать внутрь его ответ (как только он появится). Давайте напишем простенький класс, который умеет это делать.

import collections

from telegram.ext import Filters
from telegram.ext import MessageHandler
from telegram.ext import Updater


class DialogBot(object):

    def __init__(self, token, generator):
        self.updater = Updater(token=token)  # заводим апдейтера
        handler = MessageHandler(Filters.text | Filters.command, self.handle_message)
        self.updater.dispatcher.add_handler(handler)  # ставим обработчик всех текстовых сообщений
        self.handlers = collections.defaultdict(generator)  # заводим мапу "id чата -> генератор"

    def start(self):
        self.updater.start_polling()

    def handle_message(self, bot, update):
        print("Received", update.message)
        chat_id = update.message.chat_id
        if update.message.text == "/start":
            # если передана команда /start, начинаем всё с начала -- для
            # этого удаляем состояние текущего чатика, если оно есть
            self.handlers.pop(chat_id, None)
        if chat_id in self.handlers:
            # если диалог уже начат, то надо использовать .send(), чтобы
            # передать в генератор ответ пользователя
            try:
                answer = self.handlers[chat_id].send(update.message)
            except StopIteration:
                # если при этом генератор закончился -- что делать, начинаем общение с начала
                del self.handlers[chat_id]
                # (повторно вызванный, этот метод будет думать, что пользователь с нами впервые)
                return self.handle_message(bot, update)
        else:
            # диалог только начинается. defaultdict запустит новый генератор для этого
            # чатика, а мы должны будем извлечь первое сообщение с помощью .next()
            # (.send() срабатывает только после первого yield)
            answer = next(self.handlers[chat_id])
        # отправляем полученный ответ пользователю
        print("Answer: %r" % answer)
        bot.sendMessage(chat_id=chat_id, text=answer)

Что ж, осталось сочинить диалог, который мы будем отыгрывать! Давайте поговорим о Питоне.

def dialog():
    answer = yield "Здравствуйте! Меня забыли наградить именем, а как зовут вас?"
    # убираем ведущие знаки пунктуации, оставляем только 
    # первую компоненту имени, пишем её с заглавной буквы
    name = answer.text.rstrip(".!").split()[0].capitalize()
    likes_python = yield from ask_yes_or_no("Приятно познакомиться, %s. Вам нравится Питон?" % name)
    if likes_python:
        answer = yield from discuss_good_python(name)
    else:
        answer = yield from discuss_bad_python(name)


def ask_yes_or_no(question):
    """Спросить вопрос и дождаться ответа, содержащего «да» или «нет».

    Возвращает:
        bool
    """
    answer = yield question
    while not ("да" in answer.text.lower() or "нет" in answer.text.lower()):
        answer = yield "Так да или нет?"
    return "да" in answer.text.lower()


def discuss_good_python(name):
    answer = yield "Мы с вами, %s, поразительно похожи! Что вам нравится в нём больше всего?" % name
    likes_article = yield from ask_yes_or_no("Ага. А как вам, кстати, статья на Хабре? Понравилась?")
    if likes_article:
        answer = yield "Чудно!"
    else:
        answer = yield "Жалко."
    return answer


def discuss_bad_python(name):
    answer = yield "Ай-яй-яй. %s, фу таким быть! Что именно вам так не нравится?" % name
    likes_article = yield from ask_yes_or_no(
        "Ваша позиция имеет право на существование. Статья "
        "на Хабре вам, надо полагать, тоже не понравилась?")
    if likes_article:
        answer = yield "Ну и ладно."
    else:
        answer = yield "Что «нет»? «Нет, не понравилась» или «нет, понравилась»?"
        answer = yield "Спокойно, это у меня юмор такой."
    return answer


if __name__ == "__main__":
    dialog_bot = DialogBot(sys.argv[1], dialog)
    dialog_bot.start()

И это работает! Результат выглядит примерно так:

Добавляем разметку

Боты в Telegram сильны тем, что могут кидаться в своих пользователей HTML- и Markdown-разметкой; эту возможность обойти стороной нам было бы непозволительно. Чтобы понять, как послать сообщение с разметкой, давайте взглянем на описание функции Bot.sendMessage:

    def sendMessage(self, chat_id, text,
                    parse_mode=None,
                    disable_web_page_preview=None,
                    disable_notification=False,
                    reply_to_message_id=None,
                    reply_markup=None,
                    timeout=None,
                    **kwargs):
        """Use this method to send text messages.
        Args:
            chat_id (str): ...
            text (str): ...
            parse_mode (Optional[str]): Send Markdown or HTML, if you want
                Telegram apps to show bold, italic, fixed-width text or inline
                URLs in your bot's message.
            disable_web_page_preview (Optional[bool]): ...
            disable_notification (Optional[bool]): ...
            reply_to_message_id (Optional[int]): ...
            reply_markup (Optional[:class:`telegram.ReplyMarkup`]): Additional
                interface options. A JSON-serialized object for an inline
                keyboard, custom reply keyboard, instructions to hide reply
                keyboard or to force a reply from the user.
            timeout (Optional[float]): ...
        ...
        """

Ага! Достаточно передавать аргумент parse_mode="HTML" или parse_mode="Markdown". Можно было бы просто добавить это в наш вызов, но давайте сделаем чуть-чуть погибче: добавим специальные объекты, которые нужно будет yield'ить, чтобы спровоцировать использование разметки:

class Message(object):
    def __init__(self, text, **options):
        self.text = text
        self.options = options


class Markdown(Message):
    def __init__(self, text, **options):
        super(Markup, self).__init__(text, parse_mode="Markdown", **options)


class HTML(Message):
    def __init__(self, text, **options):
        super(HTML, self).__init__(text, parse_mode="HTML", **options)

Теперь отправка сообщений будет выглядеть так:

    def handle_message(self, bot, update):
        # ......
        print("Answer: %r" % answer)
        self._send_answer(bot, chat_id, answer)

    def _send_answer(self, bot, chat_id, answer):
        if isinstance(answer, str):
            answer = Message(answer)
        bot.sendMessage(chat_id=chat_id, text=answer.text, **answer.options)

Для демонстрации давайте модифицируем ask_yes_or_no():

def ask_yes_or_no(question):
    answer = yield question
    while not ("да" in answer.text.lower() or "нет" in answer.text.lower()):
        answer = yield HTML("Так <b>да</b> или <b>нет</b>?")
    return "да" in answer.text.lower()

Результат налицо:

Добавляем кнопки

Единственное, чего нам не хватает и что могло бы вполне себе пригодиться при написании диалоговых ботов — клавиатура с выбором вариантов ответа. Для создания клавиатуры нам достаточно добавить в Message.options ключ reply_markup; но давайте постараемся максимально упростить и абстрагировать наш код внутри генераторов. Здесь напрашивается решение попроще. Пусть, например, yield выдаёт не один объект, а сразу несколько; если среди них есть список или список списков со строками, например:

answer = yield (
    "Выберите цифру от 1 до 9",
    [
        ["1", "2", "3"],
        ["4", "5", "6"],
        ["7", "8", "9"],
    ]
)

, то мы считаем, что это кнопки клавиатуры, и хотим получить примерно следующий результат:



_send_answer() тогда преобразуется в нечто такое:

def _send_answer(self, bot, chat_id, answer):
        print("Sending answer %r to %s" % (answer, chat_id))
        if isinstance(answer, collections.abc.Iterable) and not isinstance(answer, str):
            # мы получили несколько объектов -- сперва каждый надо обработать
            answer = list(map(self._convert_answer_part, answer))
        else:
            # мы получили один объект -- сводим к более общей задаче
            answer = [self._convert_answer_part(answer)]
        
        # перед тем, как отправить очередное сообщение, идём вперёд в поисках
        # «довесков» -- клавиатуры там или в перспективе ещё чего-нибудь
        current_message = None
        for part in answer:
            if isinstance(part, Message):
                if current_message is not None:
                    # поскольку не все объекты исчерпаны, пусть это сообщение
                    # не вызывает звоночек (если не указано обратное)
                    options = dict(current_message.options)
                    options.setdefault("disable_notification", True)
                    bot.sendMessage(chat_id=chat_id, text=current_message.text, **options)
                current_message = part
            if isinstance(part, ReplyMarkup):
                # ага, а вот и довесок! добавляем текущему сообщению.
                # нет сообщения -- ну извините, это ошибка.
                current_message.options["reply_markup"] = part
        # надо не забыть отправить последнее встреченное сообщение.
        if current_message is not None:
            bot.sendMessage(chat_id=chat_id, text=current_message.text, **current_message.options)

    def _convert_answer_part(self, answer_part):
        if isinstance(answer_part, str):
            return Message(answer_part)
        if isinstance(answer_part, collections.abc.Iterable):
            # клавиатура?
            answer_part = list(answer_part)
            if isinstance(answer_part[0], str):
                # она! оформляем как горизонтальный ряд кнопок.
                # кстати, все наши клавиатуры одноразовые -- нам пока хватит.
                return ReplyKeyboardMarkup([answer_part], one_time_keyboard=True)
            elif isinstance(answer_part[0], collections.abc.Iterable):
                # двумерная клавиатура?
                if isinstance(answer_part[0][0], str):
                    # она!
                    return ReplyKeyboardMarkup(map(list, answer_part), one_time_keyboard=True)
        return answer_part

В качестве демонстрации поменяем ask_yes_or_no() и discuss_bad_python():

def ask_yes_or_no(question):
    """Спросить вопрос и дождаться ответа, содержащего «да» или «нет».

    Возвращает:
        bool
    """
    answer = yield (question, ["Да.", "Нет."])
    while not ("да" in answer.text.lower() or "нет" in answer.text.lower()):
        answer = yield HTML("Так <b>да</b> или <b>нет</b>?")
    return "да" in answer.text.lower()


def discuss_bad_python(name):
    answer = yield "Ай-яй-яй. %s, фу таким быть! Что именно вам так не нравится?" % name
    likes_article = yield from ask_yes_or_no(
        "Ваша позиция имеет право на существование. Статья "
        "на Хабре вам, надо полагать, тоже не понравилась?")
    if likes_article:
        answer = yield "Ну и ладно."
    else:
        answer = yield (
            "Что «нет»? «Нет, не понравилась» или «нет, понравилась»?",
            ["Нет, не понравилась!", "Нет, понравилась!"]
        )
        answer = yield "Спокойно, это у меня юмор такой."
    return answer

Результат:

Заключение

Генераторы в Питоне — мощный инструмент, и использование его в нужных ситуациях позволяет значительно сокращать и упрощать код. Посмотрите, как красиво мы, например, вынесли вопрос «да или нет» в отдельную функцию, притом оставив за ним право проводить дополнительное общение с пользователем. Так же мы могли бы вынести в отдельную функцию и вопрошание имени, и научить его уточнять у пользователя, верно ли мы его поняли, и так далее, и тому подобное. Генераторы сами хранят за нас состояние диалога, и сами умеют продолжать его с требуемого момента. Всё для нас!

Надеюсь, эта статья была кому-то полезной. Как водится, не стесняйтесь сообщать обо всех опечатках, орфографических и грамматических ошибках в личку. Весь код к статье лежит в репозитории на Github (ветка habrahabr-316666). На бота ссылку не дам и живым его держать, конечно, ближайшее время не буду, иначе хабраэффект накроет его вместе с моим компьютером. Успехов в создании своих диалоговых ботов 😉