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()