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.).
telegram(it appears we already have a twitter), but your suggestions were a lot more specific. \$\endgroup\$whatsapptag looming and thus a retagging effort that we should coordinate a bit. Thank you for your patience. \$\endgroup\$