2
\$\begingroup\$

This code review is for my Stack Exchange Marathon Bot project, a Telegram bot for playing a dumb little game involving the Stack Exchange network. Since the codebase is not that small, in this post I want to focus on only one aspect, namely the Telegram bot -related code, which uses the python-telegram-bot library. In the future I may make another post for the functionality of the game itself or other aspects.

While developing this project, I wanted to find some good design patterns for developing a Telegram bot using python-telegram-bot that I could reuse for later projects. So, the general focus for this post is to evaluate the approach I used for the design of bot-related code (but, of course, you can review whatever you feel like).

Since this still leaves a lot of code to review, I want to narrow it down even more. Telegram bots mainly work by means of /commands that the users request to execute. The general philosophy of my approach is to make altering, adding, and maintaining commands as easy as possible, as well as organising data related to the bot (e.g. session data) in an intuitive and versatile way. This will be the specific focus of this post; in the future I might make other posts focusing on other parts of the bot code.

General design

The main class in this module is SEMarathonBotSystem. It aggregates all the components form python-telegram-bot needed to run the bot, such as the Bot object, the Dispatcher, the Updater, the JobQueue, etc. It also abstracts the concept of "an instance of the Stack Exchange Marathon Bot", i.e. each instance represents the deployment of this bot (intended as a specific behaviour) to one Telegram bot "account". For example, one instance could be deployed to @BotInstanceA and another to @BotInstanceB; they would have the same behaviour, but would be different bots. In practice I only deploy it to one bot account, but I think having this abstraction makes it easier to reason about; e.g., data that is global to one bot instance would belong to the corresponding SEMarathonBotSystem instance.

Nested within this class, there is the Session class (ergo, SEMarathonBotSystem.Session). Each session is bound to a specific Telegram chat, and is the object that contains all state related to that chat (e.g. the active game instance, the configuration, etc.).

At the top level within SEMarathonBotSystem, there are stateless commands and a command to start a Session. On the other hand, Session contains stateful commands that operate on an active session.

While python-telegram-bot does provide built-in support for bot data and chat data in the form of arbitrary key-to-value dictionaries, I prefer using a direct object-oriented approach to storing data and state. So I ended up with a compromise: I add the appropriate SEMarathonBotSystem instance to the bot_data dictionary, and then work entirely on the object. Similarly, when I start a Session, I add an entry to the chat_data dictionary containing the corresponding session object. This also allows me, as you'll see below, to retrieve the correct bot system and session objects when an unbound command handler is called, so as to pass it as the self argument to the corresponding method.

Code

All code shown in this post belongs to a single module, semarathon.bot. Instead of dumping all the code here, I thought it might be easier for you if I chunk it up and explain it along the way, omitting unnecessary details or repetitive bulk. However, you can always access the full code here (pinned version) or here (most recent version).

I'm going to show fragments of code in a logical top-down, depth-first order, not necessarily the same order as in the source code.

Preamble

First, some imports and aliases, just so you know where each identifier comes from:

import datetime
import enum
import functools
import inspect
import itertools
import logging
import time
from typing import Any, Callable, Generator, List, Optional, TypeVar

import more_itertools
import telegram as tg
import telegram.ext as tge
from telegram.parsemode import ParseMode
from telegram.utils.helpers import escape_markdown as escape_md

from semarathon import marathon as mth
from semarathon.utils import Decorator, Text, coroutine, format_exception_md

# logger setup
logger = logging.getLogger(__name__)
del logging  # to avoid mistakes with code completion

# type aliases
CommandCallback = Callable[[tg.Update, tge.CallbackContext], None]
CommandCallbackMethod = Callable[["BotSession", tg.Update, tge.CallbackContext], None]
BotSessionRunnable = Callable[["BotSession"], Any]
T = TypeVar("T", CommandCallback, CommandCallbackMethod)

# other aliases
escape_mdv2 = functools.partial(escape_md, version=2)

Class skeleton

Before going into details, it might be useful to take a look at a stripped down skeleton of the SEMarathonBotSystem class to have a general idea of how it's laid out. Ellipses (...) indicate hidden or collapsed code.

class SEMarathonBotSystem:
    """
    Manages all the components of one instance of the SE Marathon Bot.

    Each instance of the SE Marathon Bot corresponds one-to-one with a Telegram bot
    username. This class allows deployment of the same abstract bot behaviour to any
    Telegram bot (i.e. bot username).
    """

    bot: tg.Bot
    updater: tge.Updater
    dispatcher: tge.Dispatcher
    job_queue: tge.JobQueue

    def __init__(self, token: str, **kwargs):
        """
        Initialize a new bot instance and bind it to a certain bot username.

        Arguments are the same as for :class:`telegram.ext.Updater` with the exception
        of use_context, which is automatically set to ``True`` (and cannot be changed).
        """

        self.updater = tge.Updater(token, use_context=True, **kwargs)
        self.bot = self.updater.bot
        self.dispatcher = self.updater.dispatcher
        self.job_queue = self.updater.job_queue

        self._setup_handlers()
        self.dispatcher.bot_data["bot_system"] = self

    ...

    class Session:
        """Represents the context of the interaction with the bot in a Telegram chat."""

        bot_system: "SEMarathonBotSystem"
        id: int
        marathon: Optional[mth.Marathon]
        operation: Optional["SEMarathonBotSystem.Session.Operation"]
        jobs: List[tge.Job]

        def __init__(self, bot_system: "SEMarathonBotSystem", chat_id: int):
            """Initialize a new session and attach it to the given chat."""
            self.bot_system = bot_system
            self.id = chat_id
            self.marathon = None
            self.operation = None
            self.jobs = []

            self.bot_system.dispatcher.chat_data[chat_id]["session"] = self

        ...

        # -------------------------- Command handlers  --------------------------

        ...

        # ------------------------------- Job callbacks  ----------------------------

        ...

        # ---------------------------- Utility methods  ----------------------------

        ...

    ...

Command handlers: the cmdhandler decorator

Perhaps the most important component in this module is the cmdhandler decorator: it decorates a function or method to make it into a command handler, while managing all the tedious details like registering the handler with the dispatcher object, adding the command to the list of commands seen by the user on the Telegram app, etc. This means adding a command to the bot is as simple as adding a new method decorated with this.

def cmdhandler(
    command: str = None,
    *,
    callback_type: _CommandCallbackType = _CommandCallbackType.SESSION_METHOD,
    register: bool = True,
    **handler_kwargs,
) -> Decorator:
    """Parametrised decorator that marks a function as a callback for a command handler.

    :param command: name of bot command to add a handler for
    :param callback_type: type of callback ("standard" top-level function, bot system
                          method, or session method)
    :param register: whether to register the command in the command list shown on
                     Telegram clients. If set to ``True``, a ``command_info`` attribute
                     is added to the callback function containing a BotCommand object,
                     and a positive integer is assigned to ``callback.register``, in
                     order of usage of this decorator.
                     Otherwise ``callback.register`` is set to ```False``.
    :param handler_kwargs: additional keyword arguments for the
                           creation of the command handler (these will be passed
                           to ``telegram.ext.dispatcher.add_handler``)

    :return: (after decoration) the decorated function, with the added
        ``command_handler`` and (optionally) ``command_info`` attributes.
    """

    def decorator(callback: T) -> T:
        command_ = command or callback.__name__

        handler = _make_command_handler(
            callback,
            command_,
            callback_type=callback_type,
            **handler_kwargs,
        )
        callback.command_handler = handler
        if register:
            callback.register = cmdhandler._counter = (
                getattr(cmdhandler, "_counter", 0) + 1
            )
            callback.command_info = tg.BotCommand(
                command_, _extract_command_description(callback)
            )
        else:
            callback.register = False

        return callback

    return decorator

Note that, by default, the decorated function's name is used as the command name, eliminating a potential source of redundancy. The command handler object itself is created in the _make_command_handler helper, shown below. This makes a decorated callback that manages things like logging, exception handling, and other bookkeeping.

def _make_command_handler(
    callback: T,
    command: str = None,
    *,
    callback_type: _CommandCallbackType = _CommandCallbackType.SESSION_METHOD,
    **handler_kwargs,
) -> tge.CommandHandler:
    """Make a command handler for the command ``command``

    Constructs a :class:`telegram.ext.CommandHandler` with a decorated version of
    the given callback. If ``command`` is not specified it defaults to the callback
    function's name. The callback is decorated with an exception handler and some
    logging business.

    :param callback: original callback for the command
    :param command: see :func:`cmdhandler`
    :param callback_type: see :func:`cmdhandler`
    :param handler_kwargs: see :func:`cmdhandler`

    :return: a command handler for the given command
    """

    command = command or callback.__name__
    callback_type = _CommandCallbackType(callback_type)

    @functools.wraps(callback)
    def decorated(update: tg.Update, context: tge.CallbackContext):
        command_info = f"/{command}@{update.effective_chat.id}"
        logger.info(f"reached {command_info}")
        try:
            bot_system = _get_bot_system(context)

            # Build arguments list:
            args = [update, context]
            if callback_type == _CommandCallbackType.SESSION_METHOD:
                args.insert(0, _get_session(context))
            elif callback_type == _CommandCallbackType.BOT_SYSTEM_METHOD:
                args.insert(0, bot_system)

            # Actual call:
            callback(*args)

            logger.info(f"served {command_info}")
        except (UsageError, ValueError, mth.SEMarathonError) as e:
            text = (
                f"{Text.load('usage-error')}\n{format_exception_md(e)}\n\n"
                f"{escape_mdv2(getattr(e, 'help_txt', 'See /info for usage info'))}"
            )
            markdown_safe_send(context.bot, update.effective_chat.id, text)
            logger.info(f"served {command_info} (with usage/algorithm error)")
        except Exception as e:
            text = f"{Text.load('internal-error')}"
            markdown_safe_send(context.bot, update.effective_chat.id, text)
            logger.exception(f"{command_info}: unexpected exception", exc_info=e)
        finally:
            logger.debug(f"exiting {command_info}")

    handler = tge.CommandHandler(command, decorated, **handler_kwargs)
    return handler

Note that exceptions are distinguished between (1) usage errors (e.g. the user providing invalid arguments to a command) and "foreseeable" exceptions (UsageError, ValueError and SEMarathonError), and (2) internal server errors (unexpected issues caused by the implementation). This allows the user to be informed accordingly.

In python-telegram-bot, all command callbacks have the same signature: callback(Update, CallbackContext). However, some of my callbacks are Session methods, some are SEMarathonBotSystem methods, etc. The callback_type parameter just determines whether to pass and what to pass as the self argument, since I've made it so command handlers can be added as SEMarathonBotSystem methods (for stateless commands), as Session methods (for commands that operate on a session), or as free functions, depending on what is most appropriate:

class _CommandCallbackType(enum.Enum):
    FREE_FUNCTION = enum.auto()
    BOT_SYSTEM_METHOD = enum.auto()
    SESSION_METHOD = enum.auto()

The relevant objects are extracted from the CallbackContext object via the _get_bot_system and _get_session methods shown below. Respectively, these use the "bot_system" and "session" entries in bot_data and chat_data, which are initialized in SEMarathonBotSystem.__init__ and SEMarathonBotSystem.Session.__init__ respectively.

def _get_bot_system(context: tge.CallbackContext) -> SEMarathonBotSystem:
    try:
        return context.bot_data["bot_system"]
    except KeyError:
        raise RuntimeError("Received an update destined for an uninitialised bot")


def _get_session(context: tge.CallbackContext) -> SEMarathonBotSystem.Session:
    try:
        return context.chat_data["session"]
    except KeyError:
        raise UsageError(
            "Session not initialized",
            help_txt="You must use /start before using other commands",
        )

Going back to the body of the cmdhanlder decorator, you can see that it creates the corresponding command handler object (and optionally a BotCommand info object), but doesn't actually register it with the dispatcher; rather, it simply adds some marker attributes to the function object. The actual registration is done later by the SEMarathonBotSystem._setup_handlers method (shown below), called during initialization, which collects all command handlers (by looking for these attributes) and registers them. (This is because cmdhandler is a free function, and it needs to be since it's a decorator, and as such it doesn't have a reference to the bot system object.)

class SEMarathonBotSystem:
    ...
    
    @classmethod
    def _collect_command_callbacks(cls):
        return [
            callback
            for name, callback in itertools.chain.from_iterable(
                inspect.getmembers(class_, inspect.isfunction)
                for class_ in (cls, cls.Session)
            )
            if hasattr(callback, "command_handler")
        ]

    def _setup_handlers(self):
        callbacks = self._collect_command_callbacks()
        for callback in callbacks:
            logger.debug(
                f"Adding command handler for {callback.command_handler.command}"
            )
            self.dispatcher.add_handler(callback.command_handler)

        cmd_list = [
            callback.command_info
            for callback in sorted(callbacks, key=lambda c: c.register)
            if callback.register
        ]
        logger.debug(
            f"Registering command list on Telegram: {[cmd.command for cmd in cmd_list]}"
        )
        self.bot.set_my_commands(cmd_list)

This method also registers all the commands on Telegram, so that the Telegram UI shows which commands are available, along with a short description for each. This information is stored in the command_info attribute of command functions, added by cmdhandler. If you look at the cmdhandler source, you'll see that the description is being extracted with the _extract_command_description helper, which, in turn, extracts it from the function's docstring:

def _extract_command_description(callback: Callable) -> str:
    doc = getattr(callback, "__doc__", None)
    if not doc:
        raise ValueError(f"No command description found for callback {callback}")
    return doc.strip().split("\n", maxsplit=1)[0]

Also, note that the final thing that is returned from the cmdhandler decoration is not the decorated version from _make_command_handler, which is used as the command handler callback, but rather it's the original function object with some extra attributes. This allows these methods to be called as usual with the same exact signature as shown in the source code from, say, an interactive Python console, without all the extra harnesses that are only useful in "automatic mode", i.e. when responding to an actual request from a Telegram user, in contrast to "manual" calls.

Command examples

Now that I've shown how the cmdhandler decorator works, let's see how I've used it.

These are two commands defined at the SEMarathonBotSystem top-level:

class SEMarathonBotSystem:
    ...
    
    # noinspection PyUnusedLocal
    @staticmethod
    @cmdhandler(callback_type=_CommandCallbackType.FREE_FUNCTION)
    def info(update: tg.Update, context: tge.CallbackContext):
        """General information about this bot and credits"""
        update.message.reply_markdown_v2(Text.load("info"))

    # noinspection PyUnusedLocal
    @cmdhandler(callback_type=_CommandCallbackType.BOT_SYSTEM_METHOD)
    def start(self, update: tg.Update, context: tge.CallbackContext) -> "Session":
        """Start a session in the current chat (start listening for commands)"""
        chat_id = update.message.chat_id
        session = SEMarathonBotSystem.Session(self, chat_id)
        update.message.reply_text(text=Text.load("start"))
        return session
    
    ...

This defines two stateless commands, /info and /start, with the description given by their docstring.

Here is a sample of command callbacks as Session methods:

class SEMarathonBotSystem:
    
    ...
    
    class Session:
        
        ...
        
        @cmdhandler()
        def new_marathon(
            self, update: tg.Update, context: tge.CallbackContext
        ) -> mth.Marathon:
            """Create a new marathon"""
            self.marathon = mth.Marathon()
            self.send_message(text=Text.load("new-marathon"))
            return self.marathon

        ...
        
        @cmdhandler()
        @marathon_method
        def status(self, update: tg.Update, context: tge.CallbackContext):
            """Show the status of the current marathon"""
            self.send_message(text=self._status_text())
        
        ...
        
        @cmdhandler()
        def stop_marathon(self, update: tg.Update, context: tge.CallbackContext):
            """Stop the marathon prematurely"""
            self.marathon.stop()
            assert not self.marathon.is_running

As an example, if we want to add a /spam command that sends the word "SPAM" times the number of times /spam has been called in the session, it's as simple as this:

        @cmdhandler()
        def spam(self, update: tg.Update, context: tge.CallbackContext):
            """Send a bunch of SPAM."""
            spam_count = self._spam_count = getattr(self, "_spam_count", 0) + 1
            self.send_message("SPAM" * spam_count)

This shows how the whole change can just happen in one place (the callback definition) rather than in several disconnected places (the callback definition, the registration of the command handler, the registration of the command info, etc.).

\$\endgroup\$
2
  • \$\begingroup\$ Note: I've removed your telegram-bot & python-telegram-bot tags. Not every project deserves its own tag. Perhaps a case can be made for the tag telegram (it appears we already have a twitter), but your suggestions were a lot more specific. \$\endgroup\$ – Mast Apr 12 at 4:22
  • \$\begingroup\$ I'll write up a proposal on meta for that tag later today. Supporting it would mean there's also a whatsapp tag looming and thus a retagging effort that we should coordinate a bit. Thank you for your patience. \$\endgroup\$ – Mast Apr 12 at 4:26
2
\$\begingroup\$

Disclaimer: I'm currently the maintainer of python-telegram-bot.


Hi.

TBH I haven't tried to understand every detail of what you presented. I'd still like to comment on the design you chose to use methods of your custom classes as callbacks. More precisely I find it rather irritating that you define methods like

class SEMarathonBotSystem:
    ...

    @cmdhandler(callback_type=_CommandCallbackType.BOT_SYSTEM_METHOD)
    def start(self, update: tg.Update, context: tge.CallbackContext) -> "Session":
        ...

but pass the self argument through some involved logic within the decorator after reading it from the context argument. IISC the necessity for this comes from the fact, that you try to do the registering of the commands in the decorators, i.e. before actually building an instance of SEMarathonBotSystem. Note, that you don't have to worry about self at all, if you work with an instance. Here is a minimal exapmle for demonstration, which I personally find much more straight forward to grasp:

from telegram import Update
from telegram.ext import Updater, CommandHandler, CallbackContext


class Foo:
    def __init__(self, dispatcher):
        self.counter = 0
        dispatcher.add_handler(CommandHandler("start", self.callback))

    def callback(self, update, _):
        self.counter += 1
        update.effective_message.reply_text(self.counter)


def main():
    updater = Updater("TOKEN")
    foo = Foo(updater.dispatcher)
    updater.start_polling()
    updater.idle()


if __name__ == '__main__':
    main()

Surely you can still use decorators to minimize duplication of the dispatcher.add_handler(CommandHandler("start", self.callback)) call, i.e. the decorators job would be to mark the class methods that are to be used as callbacks for a CommandHandler.

Of course the situation is a bit different for your Session class, as you need a session per chat and can't register a different handler for every chat in advance. Still, IMHO the choice to pass a self argument in a complicated way seems strange. In fact I think I would be less irritated, if you'd just renamed it and made it a static method:

    class Session:
        ...
        
        @statcmethod
        @cmdhandler()
        def new_marathon(
            update: tg.Update, context: tge.CallbackContext, session: 'Session'
        ) -> mth.Marathon:
            ...

However, I'd still argue that PTB brings built-in mechanisms that can take care of things like this, most importantly the context argument. That argument is designed to bring a convenient way to pass multiple additional objects to a callback without letting the signature explode (as it did in PTB<v12) and you can leverage this here in at least 2 different ways:

  1. Make your decorator do something like context.session = context.chat_data['session']
  2. Implement a subclass of CommandHandler that does exactly that within collect_additional_context

Then you can access the session as context.session. Admittedly, this has two downsides:

  1. You need to type context., but esp. with auto-completion that's neglectable IMHO
  2. Linters and type checkers will give you trouble and tell you that 'CallbackContext' has no attribute 'session'. Note that this will be resolved by PTB/#2262, which is on our roadmap and will allow you to use a custom subclass of CallbackContext.

So this is my 2cts. I'd like to mention that in my experience the "how do I structure my PTB-based bot?" is something that everyone has to find for themself and once you've done that, it's hard to get comfortable with other structures. This in addition with my "job"-related urge to promote built-in features of PTB may lead to somewhat of a bias here and in addition this is my first post on this platform :D Still, I hope that my comments are of some use for you and wish you good luck with your project!


\$\endgroup\$
1
  • \$\begingroup\$ Thanks for the feedback! I'm not going to try to hide the fact that the whole callback_type=_CommandCallbackType.BOT_SYSTEM_METHOD thing is quite ugly—I fully agree. But it would bug me even more to have something like new_marathon, which is conceptually and in essence an instance method of Session, be something other than an actual instance method! Although probably that's the lesser of the two evils. So now I just need to convince myself of that... :) In any case, this got me thinking of possible compromises. \$\endgroup\$ – Anakhand Apr 11 at 20:55

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.