contexttypesbot.pyΒΆ

  1#!/usr/bin/env python
  2# pylint: disable=unused-argument, wrong-import-position
  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
 15from typing import DefaultDict, Optional, Set
 16
 17from telegram import __version__ as TG_VER
 18
 19try:
 20    from telegram import __version_info__
 21except ImportError:
 22    __version_info__ = (0, 0, 0, 0, 0)  # type: ignore[assignment]
 23
 24if __version_info__ < (20, 0, 0, "alpha", 1):
 25    raise RuntimeError(
 26        f"This example is not compatible with your current PTB version {TG_VER}. To view the "
 27        f"{TG_VER} version of this example, "
 28        f"visit https://docs.python-telegram-bot.org/en/v{TG_VER}/examples.html"
 29    )
 30from telegram import InlineKeyboardButton, InlineKeyboardMarkup, Update
 31from telegram.constants import ParseMode
 32from telegram.ext import (
 33    Application,
 34    CallbackContext,
 35    CallbackQueryHandler,
 36    CommandHandler,
 37    ContextTypes,
 38    ExtBot,
 39    TypeHandler,
 40)
 41
 42# Enable logging
 43logging.basicConfig(
 44    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", level=logging.INFO
 45)
 46logger = logging.getLogger(__name__)
 47
 48
 49class ChatData:
 50    """Custom class for chat_data. Here we store data per message."""
 51
 52    def __init__(self) -> None:
 53        self.clicks_per_message: DefaultDict[int, int] = defaultdict(int)
 54
 55
 56# The [ExtBot, dict, ChatData, dict] is for type checkers like mypy
 57class CustomContext(CallbackContext[ExtBot, dict, ChatData, dict]):
 58    """Custom class for context."""
 59
 60    def __init__(
 61        self,
 62        application: Application,
 63        chat_id: Optional[int] = None,
 64        user_id: Optional[int] = None,
 65    ):
 66        super().__init__(application=application, chat_id=chat_id, user_id=user_id)
 67        self._message_id: Optional[int] = None
 68
 69    @property
 70    def bot_user_ids(self) -> Set[int]:
 71        """Custom shortcut to access a value stored in the bot_data dict"""
 72        return self.bot_data.setdefault("user_ids", set())
 73
 74    @property
 75    def message_clicks(self) -> Optional[int]:
 76        """Access the number of clicks for the message this context object was built for."""
 77        if self._message_id:
 78            return self.chat_data.clicks_per_message[self._message_id]
 79        return None
 80
 81    @message_clicks.setter
 82    def message_clicks(self, value: int) -> None:
 83        """Allow to change the count"""
 84        if not self._message_id:
 85            raise RuntimeError("There is no message associated with this context object.")
 86        self.chat_data.clicks_per_message[self._message_id] = value
 87
 88    @classmethod
 89    def from_update(cls, update: object, application: "Application") -> "CustomContext":
 90        """Override from_update to set _message_id."""
 91        # Make sure to call super()
 92        context = super().from_update(update, application)
 93
 94        if context.chat_data and isinstance(update, Update) and update.effective_message:
 95            # pylint: disable=protected-access
 96            context._message_id = update.effective_message.message_id
 97
 98        # Remember to return the object
 99        return context
100
101
102async def start(update: Update, context: CustomContext) -> None:
103    """Display a message with a button."""
104    await update.message.reply_html(
105        "This button was clicked <i>0</i> times.",
106        reply_markup=InlineKeyboardMarkup.from_button(
107            InlineKeyboardButton(text="Click me!", callback_data="button")
108        ),
109    )
110
111
112async def count_click(update: Update, context: CustomContext) -> None:
113    """Update the click count for the message."""
114    context.message_clicks += 1
115    await update.callback_query.answer()
116    await update.effective_message.edit_text(
117        f"This button was clicked <i>{context.message_clicks}</i> times.",
118        reply_markup=InlineKeyboardMarkup.from_button(
119            InlineKeyboardButton(text="Click me!", callback_data="button")
120        ),
121        parse_mode=ParseMode.HTML,
122    )
123
124
125async def print_users(update: Update, context: CustomContext) -> None:
126    """Show which users have been using this bot."""
127    await update.message.reply_text(
128        "The following user IDs have used this bot: "
129        f'{", ".join(map(str, context.bot_user_ids))}'
130    )
131
132
133async def track_users(update: Update, context: CustomContext) -> None:
134    """Store the user id of the incoming update, if any."""
135    if update.effective_user:
136        context.bot_user_ids.add(update.effective_user.id)
137
138
139def main() -> None:
140    """Run the bot."""
141    context_types = ContextTypes(context=CustomContext, chat_data=ChatData)
142    application = Application.builder().token("TOKEN").context_types(context_types).build()
143
144    # run track_users in its own group to not interfere with the user handlers
145    application.add_handler(TypeHandler(Update, track_users), group=-1)
146    application.add_handler(CommandHandler("start", start))
147    application.add_handler(CallbackQueryHandler(count_click))
148    application.add_handler(CommandHandler("print_users", print_users))
149
150    application.run_polling()
151
152
153if __name__ == "__main__":
154    main()