Skip to content
Back to top
View this page

# `contexttypesbot.py`

      1#!/usr/bin/env python
      2# pylint: disable=unused-argument
      3# This program is dedicated to the public domain under the CC0 license.
      4
      5"""
      6Simple Bot to showcase `telegram.ext.ContextTypes`.
      7
      8Usage:
      9Press Ctrl-C on the command line or send a signal to the process to stop the
     10bot.
     11"""
     12
     13import logging
     14from collections import defaultdict
     15
     16from telegram import InlineKeyboardButton, InlineKeyboardMarkup, Update
     17from telegram.constants import ParseMode
     18from telegram.ext import (
     19    Application,
     20    CallbackContext,
     21    CallbackQueryHandler,
     22    CommandHandler,
     23    ContextTypes,
     24    ExtBot,
     25    TypeHandler,
     26)
     27
     28# Enable logging
     29logging.basicConfig(
     30    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", level=logging.INFO
     31)
     32# set higher logging level for httpx to avoid all GET and POST requests being logged
     33logging.getLogger("httpx").setLevel(logging.WARNING)
     34
     35logger = logging.getLogger(__name__)
     36
     37
     38class ChatData:
     39    """Custom class for chat_data. Here we store data per message."""
     40
     41    def __init__(self) -> None:
     42        self.clicks_per_message: defaultdict[int, int] = defaultdict(int)
     43
     44
     45# The [ExtBot, dict, ChatData, dict] is for type checkers like mypy
     46class CustomContext(CallbackContext[ExtBot, dict, ChatData, dict]):
     47    """Custom class for context."""
     48
     49    def __init__(
     50        self,
     51        application: Application,
     52        chat_id: int | None = None,
     53        user_id: int | None = None,
     54    ):
     55        super().__init__(application=application, chat_id=chat_id, user_id=user_id)
     56        self._message_id: int | None = None
     57
     58    @property
     59    def bot_user_ids(self) -> set[int]:
     60        """Custom shortcut to access a value stored in the bot_data dict"""
     61        return self.bot_data.setdefault("user_ids", set())
     62
     63    @property
     64    def message_clicks(self) -> int | None:
     65        """Access the number of clicks for the message this context object was built for."""
     66        if self._message_id:
     67            return self.chat_data.clicks_per_message[self._message_id]
     68        return None
     69
     70    @message_clicks.setter
     71    def message_clicks(self, value: int) -> None:
     72        """Allow to change the count"""
     73        if not self._message_id:
     74            raise RuntimeError("There is no message associated with this context object.")
     75        self.chat_data.clicks_per_message[self._message_id] = value
     76
     77    @classmethod
     78    def from_update(cls, update: object, application: "Application") -> "CustomContext":
     79        """Override from_update to set _message_id."""
     80        # Make sure to call super()
     81        context = super().from_update(update, application)
     82
     83        if context.chat_data and isinstance(update, Update) and update.effective_message:
     84            # pylint: disable=protected-access
     85            context._message_id = update.effective_message.message_id
     86
     87        # Remember to return the object
     88        return context
     89
     90
     91async def start(update: Update, context: CustomContext) -> None:
     92    """Display a message with a button."""
     93    await update.message.reply_html(
     94        "This button was clicked <i>0</i> times.",
     95        reply_markup=InlineKeyboardMarkup.from_button(
     96            InlineKeyboardButton(text="Click me!", callback_data="button")
     97        ),
     98    )
     99
    100
    101async def count_click(update: Update, context: CustomContext) -> None:
    102    """Update the click count for the message."""
    103    context.message_clicks += 1
    104    await update.callback_query.answer()
    105    await update.effective_message.edit_text(
    106        f"This button was clicked <i>{context.message_clicks}</i> times.",
    107        reply_markup=InlineKeyboardMarkup.from_button(
    108            InlineKeyboardButton(text="Click me!", callback_data="button")
    109        ),
    110        parse_mode=ParseMode.HTML,
    111    )
    112
    113
    114async def print_users(update: Update, context: CustomContext) -> None:
    115    """Show which users have been using this bot."""
    116    await update.message.reply_text(
    117        f"The following user IDs have used this bot: {', '.join(map(str, context.bot_user_ids))}"
    118    )
    119
    120
    121async def track_users(update: Update, context: CustomContext) -> None:
    122    """Store the user id of the incoming update, if any."""
    123    if update.effective_user:
    124        context.bot_user_ids.add(update.effective_user.id)
    125
    126
    127def main() -> None:
    128    """Run the bot."""
    129    context_types = ContextTypes(context=CustomContext, chat_data=ChatData)
    130    application = Application.builder().token("TOKEN").context_types(context_types).build()
    131
    132    # run track_users in its own group to not interfere with the user handlers
    133    application.add_handler(TypeHandler(Update, track_users), group=-1)
    134    application.add_handler(CommandHandler("start", start))
    135    application.add_handler(CallbackQueryHandler(count_click))
    136    application.add_handler(CommandHandler("print_users", print_users))
    137
    138    application.run_polling(allowed_updates=Update.ALL_TYPES)
    139
    140
    141if __name__ == "__main__":
    142    main()