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