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__(self, application: Application, chat_id: int = None, user_id: int = None):
 61        super().__init__(application=application, chat_id=chat_id, user_id=user_id)
 62        self._message_id: Optional[int] = None
 63
 64    @property
 65    def bot_user_ids(self) -> Set[int]:
 66        """Custom shortcut to access a value stored in the bot_data dict"""
 67        return self.bot_data.setdefault("user_ids", set())
 68
 69    @property
 70    def message_clicks(self) -> Optional[int]:
 71        """Access the number of clicks for the message this context object was built for."""
 72        if self._message_id:
 73            return self.chat_data.clicks_per_message[self._message_id]
 74        return None
 75
 76    @message_clicks.setter
 77    def message_clicks(self, value: int) -> None:
 78        """Allow to change the count"""
 79        if not self._message_id:
 80            raise RuntimeError("There is no message associated with this context object.")
 81        self.chat_data.clicks_per_message[self._message_id] = value
 82
 83    @classmethod
 84    def from_update(cls, update: object, application: "Application") -> "CustomContext":
 85        """Override from_update to set _message_id."""
 86        # Make sure to call super()
 87        context = super().from_update(update, application)
 88
 89        if context.chat_data and isinstance(update, Update) and update.effective_message:
 90            # pylint: disable=protected-access
 91            context._message_id = update.effective_message.message_id
 92
 93        # Remember to return the object
 94        return context
 95
 96
 97async def start(update: Update, context: CustomContext) -> None:
 98    """Display a message with a button."""
 99    await update.message.reply_html(
100        "This button was clicked <i>0</i> times.",
101        reply_markup=InlineKeyboardMarkup.from_button(
102            InlineKeyboardButton(text="Click me!", callback_data="button")
103        ),
104    )
105
106
107async def count_click(update: Update, context: CustomContext) -> None:
108    """Update the click count for the message."""
109    context.message_clicks += 1
110    await update.callback_query.answer()
111    await update.effective_message.edit_text(
112        f"This button was clicked <i>{context.message_clicks}</i> times.",
113        reply_markup=InlineKeyboardMarkup.from_button(
114            InlineKeyboardButton(text="Click me!", callback_data="button")
115        ),
116        parse_mode=ParseMode.HTML,
117    )
118
119
120async def print_users(update: Update, context: CustomContext) -> None:
121    """Show which users have been using this bot."""
122    await update.message.reply_text(
123        "The following user IDs have used this bot: "
124        f'{", ".join(map(str, context.bot_user_ids))}'
125    )
126
127
128async def track_users(update: Update, context: CustomContext) -> None:
129    """Store the user id of the incoming update, if any."""
130    if update.effective_user:
131        context.bot_user_ids.add(update.effective_user.id)
132
133
134def main() -> None:
135    """Run the bot."""
136    context_types = ContextTypes(context=CustomContext, chat_data=ChatData)
137    application = Application.builder().token("TOKEN").context_types(context_types).build()
138
139    # run track_users in its own group to not interfere with the user handlers
140    application.add_handler(TypeHandler(Update, track_users), group=-1)
141    application.add_handler(CommandHandler("start", start))
142    application.add_handler(CallbackQueryHandler(count_click))
143    application.add_handler(CommandHandler("print_users", print_users))
144
145    application.run_polling()
146
147
148if __name__ == "__main__":
149    main()