Skip to content

Commit 10bdf82

Browse files
jh0kerjh0ker
authored
Add pass_user_data and pass_chat_data to Handler (python-telegram-bot#436)
* initial commit for user_data * add chat_data and use defaultdict * fix chat_data copy-paste error * add test for user_data and chat_data * fix case where chat is None * remove braces from import line
1 parent 45936c9 commit 10bdf82

File tree

12 files changed

+357
-45
lines changed

12 files changed

+357
-45
lines changed

‎examples/conversationbot2.py‎

Lines changed: 152 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,152 @@
1+
#!/usr/bin/env python
2+
# -*- coding: utf-8 -*-
3+
#
4+
# Simple Bot to reply to Telegram messages
5+
# This program is dedicated to the public domain under the CC0 license.
6+
"""
7+
This Bot uses the Updater class to handle the bot.
8+
9+
First, a few callback functions are defined. Then, those functions are passed to
10+
the Dispatcher and registered at their respective places.
11+
Then, the bot is started and runs until we press Ctrl-C on the command line.
12+
13+
Usage:
14+
Example of a bot-user conversation using ConversationHandler.
15+
Send /start to initiate the conversation.
16+
Press Ctrl-C on the command line or send a signal to the process to stop the
17+
bot.
18+
"""
19+
20+
from telegram import ReplyKeyboardMarkup
21+
from telegram.ext import (Updater, CommandHandler, MessageHandler, Filters, RegexHandler,
22+
ConversationHandler)
23+
24+
import logging
25+
26+
# Enable logging
27+
logging.basicConfig(format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
28+
level=logging.INFO)
29+
30+
logger = logging.getLogger(__name__)
31+
32+
CHOOSING, TYPING_REPLY, TYPING_CHOICE = range(3)
33+
34+
reply_keyboard = [['Age', 'Favourite colour'],
35+
['Number of siblings', 'Something else...'],
36+
['Done']]
37+
markup = ReplyKeyboardMarkup(reply_keyboard, one_time_keyboard=True)
38+
39+
40+
def facts_to_str(user_data):
41+
facts = list()
42+
43+
for key, value in user_data.items():
44+
facts.append('%s - %s' % (key, value))
45+
46+
return "\n".join(facts).join(['\n', '\n'])
47+
48+
49+
def start(bot, update):
50+
update.message.reply_text(
51+
"Hi! My name is Doctor Botter. I will hold a more complex conversation with you. "
52+
"Why don't you tell me something about yourself?",
53+
reply_markup=markup)
54+
55+
return CHOOSING
56+
57+
58+
def regular_choice(bot, update, user_data):
59+
text = update.message.text
60+
user_data['choice'] = text
61+
update.message.reply_text('Your %s? Yes, I would love to hear about that!' % text.lower())
62+
63+
return TYPING_REPLY
64+
65+
66+
def custom_choice(bot, update):
67+
update.message.reply_text('Alright, please send me the category first, '
68+
'for example "Most impressive skill"')
69+
70+
return TYPING_CHOICE
71+
72+
73+
def received_information(bot, update, user_data):
74+
text = update.message.text
75+
category = user_data['choice']
76+
user_data[category] = text
77+
del user_data['choice']
78+
79+
update.message.reply_text("Neat! Just so you know, this is what you already told me:"
80+
"%s"
81+
"You can tell me more, or change your opinion on something."
82+
% facts_to_str(user_data),
83+
reply_markup=markup)
84+
85+
return CHOOSING
86+
87+
88+
def done(bot, update, user_data):
89+
if 'choice' in user_data:
90+
del user_data['choice']
91+
92+
update.message.reply_text("I learned these facts about you:"
93+
"%s"
94+
"Until next time!" % facts_to_str(user_data))
95+
96+
user_data.clear()
97+
return ConversationHandler.END
98+
99+
100+
def error(bot, update, error):
101+
logger.warn('Update "%s" caused error "%s"' % (update, error))
102+
103+
104+
def main():
105+
# Create the Updater and pass it your bot's token.
106+
updater = Updater("TOKEN")
107+
108+
# Get the dispatcher to register handlers
109+
dp = updater.dispatcher
110+
111+
# Add conversation handler with the states GENDER, PHOTO, LOCATION and BIO
112+
conv_handler = ConversationHandler(
113+
entry_points=[CommandHandler('start', start)],
114+
115+
states={
116+
CHOOSING: [RegexHandler('^(Age|Favourite colour|Number of siblings)$',
117+
regular_choice,
118+
pass_user_data=True),
119+
RegexHandler('^Something else...$',
120+
custom_choice),
121+
],
122+
123+
TYPING_CHOICE: [MessageHandler([Filters.text],
124+
regular_choice,
125+
pass_user_data=True),
126+
],
127+
128+
TYPING_REPLY: [MessageHandler([Filters.text],
129+
received_information,
130+
pass_user_data=True),
131+
],
132+
},
133+
134+
fallbacks=[RegexHandler('^Done$', done, pass_user_data=True)]
135+
)
136+
137+
dp.add_handler(conv_handler)
138+
139+
# log all errors
140+
dp.add_error_handler(error)
141+
142+
# Start the Bot
143+
updater.start_polling()
144+
145+
# Run the bot until the you presses Ctrl-C or the process receives SIGINT,
146+
# SIGTERM or SIGABRT. This should be used most of the time, since
147+
# start_polling() is non-blocking and will stop the bot gracefully.
148+
updater.idle()
149+
150+
151+
if __name__ == '__main__':
152+
main()

‎telegram/ext/callbackqueryhandler.py‎

Lines changed: 17 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,14 @@ class CallbackQueryHandler(Handler):
5252
pass_groupdict (optional[bool]): If the callback should be passed the
5353
result of ``re.match(pattern, data).groupdict()`` as a keyword
5454
argument called ``groupdict``. Default is ``False``
55+
pass_user_data (optional[bool]): If set to ``True``, a keyword argument called
56+
``user_data`` will be passed to the callback function. It will be a ``dict`` you
57+
can use to keep any data related to the user that sent the update. For each update of
58+
the same user, it will be the same ``dict``. Default is ``False``.
59+
pass_chat_data (optional[bool]): If set to ``True``, a keyword argument called
60+
``chat_data`` will be passed to the callback function. It will be a ``dict`` you
61+
can use to keep any data related to the chat that the update was sent in.
62+
For each update in the same chat, it will be the same ``dict``. Default is ``False``.
5563
"""
5664

5765
def __init__(self,
@@ -60,9 +68,15 @@ def __init__(self,
6068
pass_job_queue=False,
6169
pattern=None,
6270
pass_groups=False,
63-
pass_groupdict=False):
71+
pass_groupdict=False,
72+
pass_user_data=False,
73+
pass_chat_data=False):
6474
super(CallbackQueryHandler, self).__init__(
65-
callback, pass_update_queue=pass_update_queue, pass_job_queue=pass_job_queue)
75+
callback,
76+
pass_update_queue=pass_update_queue,
77+
pass_job_queue=pass_job_queue,
78+
pass_user_data=pass_user_data,
79+
pass_chat_data=pass_chat_data)
6680

6781
if isinstance(pattern, string_types):
6882
pattern = re.compile(pattern)
@@ -81,7 +95,7 @@ def check_update(self, update):
8195
return True
8296

8397
def handle_update(self, update, dispatcher):
84-
optional_args = self.collect_optional_args(dispatcher)
98+
optional_args = self.collect_optional_args(dispatcher, update)
8599
if self.pattern:
86100
match = re.match(self.pattern, update.callback_query.data)
87101

‎telegram/ext/choseninlineresulthandler.py‎

Lines changed: 20 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -40,17 +40,34 @@ class ChosenInlineResultHandler(Handler):
4040
``job_queue`` will be passed to the callback function. It will be a ``JobQueue``
4141
instance created by the ``Updater`` which can be used to schedule new jobs.
4242
Default is ``False``.
43+
pass_user_data (optional[bool]): If set to ``True``, a keyword argument called
44+
``user_data`` will be passed to the callback function. It will be a ``dict`` you
45+
can use to keep any data related to the user that sent the update. For each update of
46+
the same user, it will be the same ``dict``. Default is ``False``.
47+
pass_chat_data (optional[bool]): If set to ``True``, a keyword argument called
48+
``chat_data`` will be passed to the callback function. It will be a ``dict`` you
49+
can use to keep any data related to the chat that the update was sent in.
50+
For each update in the same chat, it will be the same ``dict``. Default is ``False``.
4351
"""
4452

45-
def __init__(self, callback, pass_update_queue=False, pass_job_queue=False):
53+
def __init__(self,
54+
callback,
55+
pass_update_queue=False,
56+
pass_job_queue=False,
57+
pass_user_data=False,
58+
pass_chat_data=False):
4659
super(ChosenInlineResultHandler, self).__init__(
47-
callback, pass_update_queue=pass_update_queue, pass_job_queue=pass_job_queue)
60+
callback,
61+
pass_update_queue=pass_update_queue,
62+
pass_job_queue=pass_job_queue,
63+
pass_user_data=pass_user_data,
64+
pass_chat_data=pass_chat_data)
4865

4966
def check_update(self, update):
5067
return isinstance(update, Update) and update.chosen_inline_result
5168

5269
def handle_update(self, update, dispatcher):
53-
optional_args = self.collect_optional_args(dispatcher)
70+
optional_args = self.collect_optional_args(dispatcher, update)
5471

5572
return self.callback(dispatcher.bot, update, **optional_args)
5673

‎telegram/ext/commandhandler.py‎

Lines changed: 17 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,14 @@ class CommandHandler(Handler):
4949
``job_queue`` will be passed to the callback function. It will be a ``JobQueue``
5050
instance created by the ``Updater`` which can be used to schedule new jobs.
5151
Default is ``False``.
52+
pass_user_data (optional[bool]): If set to ``True``, a keyword argument called
53+
``user_data`` will be passed to the callback function. It will be a ``dict`` you
54+
can use to keep any data related to the user that sent the update. For each update of
55+
the same user, it will be the same ``dict``. Default is ``False``.
56+
pass_chat_data (optional[bool]): If set to ``True``, a keyword argument called
57+
``chat_data`` will be passed to the callback function. It will be a ``dict`` you
58+
can use to keep any data related to the chat that the update was sent in.
59+
For each update in the same chat, it will be the same ``dict``. Default is ``False``.
5260
"""
5361

5462
def __init__(self,
@@ -57,9 +65,15 @@ def __init__(self,
5765
allow_edited=False,
5866
pass_args=False,
5967
pass_update_queue=False,
60-
pass_job_queue=False):
68+
pass_job_queue=False,
69+
pass_user_data=False,
70+
pass_chat_data=False):
6171
super(CommandHandler, self).__init__(
62-
callback, pass_update_queue=pass_update_queue, pass_job_queue=pass_job_queue)
72+
callback,
73+
pass_update_queue=pass_update_queue,
74+
pass_job_queue=pass_job_queue,
75+
pass_user_data=pass_user_data,
76+
pass_chat_data=pass_chat_data)
6377
self.command = command
6478
self.allow_edited = allow_edited
6579
self.pass_args = pass_args
@@ -76,7 +90,7 @@ def check_update(self, update):
7690
return False
7791

7892
def handle_update(self, update, dispatcher):
79-
optional_args = self.collect_optional_args(dispatcher)
93+
optional_args = self.collect_optional_args(dispatcher, update)
8094

8195
message = update.message or update.edited_message
8296

‎telegram/ext/conversationhandler.py‎

Lines changed: 2 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222

2323
from telegram import Update
2424
from telegram.ext import Handler
25+
from telegram.utils.helpers import extract_chat_and_user
2526
from telegram.utils.promise import Promise
2627

2728

@@ -115,29 +116,7 @@ def check_update(self, update):
115116
if not isinstance(update, Update):
116117
return False
117118

118-
user = None
119-
chat = None
120-
121-
if update.message:
122-
user = update.message.from_user
123-
chat = update.message.chat
124-
125-
elif update.edited_message:
126-
user = update.edited_message.from_user
127-
chat = update.edited_message.chat
128-
129-
elif update.inline_query:
130-
user = update.inline_query.from_user
131-
132-
elif update.chosen_inline_result:
133-
user = update.chosen_inline_result.from_user
134-
135-
elif update.callback_query:
136-
user = update.callback_query.from_user
137-
chat = update.callback_query.message.chat if update.callback_query.message else None
138-
139-
else:
140-
return False
119+
chat, user = extract_chat_and_user(update)
141120

142121
key = (chat.id, user.id) if chat else (None, user.id)
143122
state = self.conversations.get(key)

‎telegram/ext/dispatcher.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
from threading import Thread, Lock, Event, current_thread, BoundedSemaphore
2525
from time import sleep
2626
from uuid import uuid4
27+
from collections import defaultdict
2728

2829
from queue import Queue, Empty
2930

@@ -86,6 +87,11 @@ def __init__(self, bot, update_queue, workers=4, exception_event=None, job_queue
8687
self.job_queue = job_queue
8788
self.workers = workers
8889

90+
self.user_data = defaultdict(dict)
91+
""":type: dict[int, dict]"""
92+
self.chat_data = defaultdict(dict)
93+
""":type: dict[int, dict]"""
94+
8995
self.handlers = {}
9096
""":type: dict[int, list[Handler]"""
9197
self.groups = []

‎telegram/ext/handler.py‎

Lines changed: 27 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
Dispatcher """
2121

2222
from telegram.utils.deprecate import deprecate
23+
from telegram.utils.helpers import extract_chat_and_user
2324

2425

2526
class Handler(object):
@@ -39,12 +40,27 @@ class Handler(object):
3940
``job_queue`` will be passed to the callback function. It will be a ``JobQueue``
4041
instance created by the ``Updater`` which can be used to schedule new jobs.
4142
Default is ``False``.
43+
pass_user_data (optional[bool]): If set to ``True``, a keyword argument called
44+
``user_data`` will be passed to the callback function. It will be a ``dict`` you
45+
can use to keep any data related to the user that sent the update. For each update of
46+
the same user, it will be the same ``dict``. Default is ``False``.
47+
pass_chat_data (optional[bool]): If set to ``True``, a keyword argument called
48+
``chat_data`` will be passed to the callback function. It will be a ``dict`` you
49+
can use to keep any data related to the chat that the update was sent in.
50+
For each update in the same chat, it will be the same ``dict``. Default is ``False``.
4251
"""
4352

44-
def __init__(self, callback, pass_update_queue=False, pass_job_queue=False):
53+
def __init__(self,
54+
callback,
55+
pass_update_queue=False,
56+
pass_job_queue=False,
57+
pass_user_data=False,
58+
pass_chat_data=False):
4559
self.callback = callback
4660
self.pass_update_queue = pass_update_queue
4761
self.pass_job_queue = pass_job_queue
62+
self.pass_user_data = pass_user_data
63+
self.pass_chat_data = pass_chat_data
4864

4965
def check_update(self, update):
5066
"""
@@ -74,7 +90,7 @@ def handle_update(self, update, dispatcher):
7490
"""
7591
raise NotImplementedError
7692

77-
def collect_optional_args(self, dispatcher):
93+
def collect_optional_args(self, dispatcher, update=None):
7894
"""
7995
Prepares the optional arguments that are the same for all types of
8096
handlers
@@ -83,10 +99,19 @@ def collect_optional_args(self, dispatcher):
8399
dispatcher (Dispatcher):
84100
"""
85101
optional_args = dict()
102+
86103
if self.pass_update_queue:
87104
optional_args['update_queue'] = dispatcher.update_queue
88105
if self.pass_job_queue:
89106
optional_args['job_queue'] = dispatcher.job_queue
107+
if self.pass_user_data or self.pass_chat_data:
108+
chat, user = extract_chat_and_user(update)
109+
110+
if self.pass_user_data:
111+
optional_args['user_data'] = dispatcher.user_data[user.id]
112+
113+
if self.pass_chat_data:
114+
optional_args['chat_data'] = dispatcher.chat_data[chat.id if chat else None]
90115

91116
return optional_args
92117

0 commit comments

Comments
 (0)