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