Add ToggleButton and RadioButton and remove the need for chat/user data

jsmnbom · jsmnbom · commit bf0a9e8df03b · 2017-05-28T13:35:00.000+02:00
Also rename the name parameter of buttons to uuid to better detail what it does, and allow the toggle and radio buttons to have names seperate from their uuid
Pass update around (unfortunately many properties aren't properties anymore)
Add a post_init method to buttons that are called after they've gotten assigned their parent_menu attribute
diff --git a/examples/menubot.py b/examples/menubot.py
@@ -8,7 +8,7 @@
 
 import telegram.ext
 from telegram.ext import CommandHandler
-from telegram.ext.menu import Menu, MenuHandler, Button, BackButton
+from telegram.ext.menu import Menu, MenuHandler, Button, BackButton, ToggleButton, RadioButton
 
 logging.basicConfig(format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
                     level=logging.DEBUG)
@@ -28,6 +28,11 @@ def start(bot, update):
     update.effective_message.reply_text('Hello!')
 
 
+def submit(bot, update, menu_data):
+    update.callback_query.answer()
+    update.effective_message.reply_text(str(menu_data))
+
+
 class MainMenu(Menu):
     def text(self, update):
         return 'Hello {}. This is the main menu!'.format(update.effective_user.first_name)
@@ -44,7 +49,7 @@ class SubMenu1(Menu):
 
     def buttons(self):
         return [
-            # [ToggleButton('Toggleable', 'on'), ToggleButton('Toggleable', 'on')],
+            [Button('URL', url='https://google.com')],
             [Button('Recursion', menu=sub_menu1), Button('Other menu!', menu=sub_menu2)],
             [BackButton('Back')]
         ]
@@ -54,9 +59,23 @@ class SubMenu2(Menu):
     text = 'This is sub menu 2'
 
     buttons = [
-            [Button('Start', start), Button('URL', url='https://google.com')],
-            # [RadioButton('Option1', 1, group=1), RadioButton('Option2', 2, group=1)],
-            [BackButton('Back')]
+            [
+                ToggleButton('test', 'Test'),
+                ToggleButton('count', states=((1, '1'), (2, '2'), (3, '3')), default=2)
+            ],
+            [
+                RadioButton('options', 1, 'Option 1'),
+                RadioButton('options', 2, 'Option 2', enabled=True),
+                RadioButton('options', 3, 'Option 3')
+            ],
+            [
+                RadioButton('custom', 1, ('[ ] Custom', '[x] Custom')),
+                RadioButton('custom', 2, ('[ ] Custom 2', '[x] Custom 2'), enabled=True)
+            ],
+            [
+                Button('Submit', submit, pass_menu_data=True),
+                BackButton('Back')
+            ]
         ]
 
 
@@ -65,7 +84,7 @@ class SubMenu2(Menu):
 sub_menu2 = SubMenu2()
 
 dp.add_handler(MenuHandler(main_menu))
-dp.add_handler(CommandHandler('menu', main_menu.start, pass_user_data=True, pass_chat_data=True))
+dp.add_handler(CommandHandler('menu', main_menu.start))
 
 # Or maybe?
 # dp.add_handler(MenuHandler(MainMenu, entry=CommandHandler('menu'))
diff --git a/telegram/ext/menu.py b/telegram/ext/menu.py
@@ -18,16 +18,19 @@
 # along with this program.  If not, see [http://www.gnu.org/licenses/].
 """This module contains objects helps with creating menus for telegram bots."""
 
-# pylint: disable=not-callable
-
-import uuid
-from collections import defaultdict
+from uuid import uuid4
+from collections import defaultdict, OrderedDict
 from itertools import chain
 
 from telegram import InlineKeyboardButton, InlineKeyboardMarkup
 from telegram.error import BadRequest
 from telegram.ext import Handler
 
+try:
+    str_type = basestring  # noqa
+except NameError:
+    str_type = str
+
 
 def id_from_update(update):
     if update.callback_query.message:
@@ -39,50 +42,51 @@ class Menu(object):
     _buttons = None
     text = ''
     buttons = None
-    data = {}
+    data = None
+    default_data = None
     root_menu = None  # populated in menuhandler
     stack = None  # Only used in root menu assigned in menuhandler
 
-    def callback(self, bot, update, user_data, chat_data):
-        self.root_menu.stack[id_from_update(update)].append(self)
+    def callback(self, bot, update, add_to_stack=True):
+        if add_to_stack:
+            self.root_menu.stack[id_from_update(update)].append(self)
+
         try:
             return update.callback_query.edit_message_text(self.get_text(update),
-                                                           reply_markup=self.keyboard(user_data,
-                                                                                      chat_data))
+                                                           reply_markup=self.keyboard(update))
         except BadRequest as e:
             if 'Message is not modified' in e.message:
                 update.callback_query.answer()
             else:
                 raise
 
-    def start(self, bot, update, user_data=None, chat_data=None):
+    def start(self, bot, update):
         # user_ and chat_data is only needed if we wanna do stuff that need state (ie.
         # ToggleButtons)
-        return update.message.reply_text(self.get_text(update), reply_markup=self.keyboard(
-            user_data, chat_data))
+        return update.message.reply_text(self.get_text(update), reply_markup=self.keyboard(update))
 
-    def keyboard(self, user_data, chat_data):
-        return InlineKeyboardMarkup([[x.keyboard_button(user_data, chat_data) for x in y] for y in
+    def keyboard(self, update):
+        return InlineKeyboardMarkup([[x.keyboard_button(update) for x in y] for y in
                                      self.get_buttons()])  # noqa
 
     def get_buttons(self):
         if self._buttons is None:
             if callable(self.buttons):
-                self._buttons = self.buttons()  # noqa
+                self._buttons = self.buttons()  # noqa pylint: disable=not-callable
             else:
                 self._buttons = self.buttons
         return self._buttons
 
     def get_text(self, update):
         if callable(self.text):
-            return self.text(update)  # noqa
+            return self.text(update)  # noqa pylint: disable=not-callable
         else:
             return self.text
 
 
 class Button(Handler):
     def __init__(self,
-                 text,
+                 text=None,
                  callback=None,
                  menu=None,
                  url=None,
@@ -94,16 +98,15 @@ def __init__(self,
                  pass_job_queue=False,
                  pass_user_data=False,
                  pass_chat_data=False,
-                 name=None):
+                 pass_menu_data=False,
+                 uuid=None):
         self._text = text
 
         if callback is not None and menu is not None:
             raise RuntimeError
         self.callback = callback
         if menu is not None:
             self.callback = menu.callback
-            pass_user_data = True
-            pass_chat_data = True
         self.menu = menu
 
         self.url = url
@@ -112,12 +115,12 @@ def __init__(self,
         self.switch_inline_query_current_chat = switch_inline_query_current_chat
         self.callback_game = callback_game
 
-        self.name = name
-        if self.name is None:
-            self.name = str(uuid.uuid4())
+        self.uuid = uuid
+        if self.uuid is None:
+            self.uuid = str(uuid4())
+        self.pass_menu_data = pass_menu_data
 
         self.parent_menu = None
-        self.root_menu = None
 
         super(Button, self).__init__(
             pass_update_queue=pass_update_queue,
@@ -129,44 +132,131 @@ def check_update(self, update):
         # Since it's not registered using Dispatcher.add_handler this really doesn't matter
         return None
 
+    def collect_optional_args(self, dispatcher, update=None):
+        optional_args = super(Button, self).collect_optional_args(dispatcher, update)
+
+        if self.pass_menu_data:
+            menu_data = self.parent_menu.root_menu.default_data.copy()
+            menu_data.update(self.parent_menu.root_menu.data[id_from_update(update)])
+            optional_args['menu_data'] = menu_data
+
+        return optional_args
+
     def handle_update(self, update, dispatcher):
         optional_args = self.collect_optional_args(dispatcher, update)
 
         return self.callback(dispatcher.bot, update, **optional_args)
 
-    def keyboard_button(self, user_data, chat_data):
-        return InlineKeyboardButton(self.text(user_data, chat_data), self.url, self.name,
+    def keyboard_button(self, update):
+        return InlineKeyboardButton(self.text(update), self.url, self.uuid,
                                     self.switch_inline_query,
                                     self.switch_inline_query_current_chat,
                                     self.callback_game)
 
-    def text(self, user_data, chat_data):
+    def text(self, update):
         return self._text
 
+    def post_init(self):
+        pass
+
 
 class BackButton(Button):
-    def __init__(self, text, name=None):
-        super(BackButton, self).__init__(text, callback=self._callback,
-                                         pass_user_data=True, pass_chat_data=True, name=name)
+    def __init__(self, text, uuid=None):
+        super(BackButton, self).__init__(text, callback=self._callback, uuid=uuid)
 
-    def _callback(self, bot, update, user_data, chat_data):
+    def _callback(self, bot, update):
         stack = self.parent_menu.root_menu.stack[id_from_update(update)]
         try:
             stack.pop()
             last_menu = stack.pop()
         except IndexError:
             last_menu = self.parent_menu.root_menu
-        last_menu.callback(bot, update, user_data, chat_data)
+        last_menu.callback(bot, update)
+
+
+class ToggleButton(Button):
+    def __init__(self, name, text=None, states=None, default=None, uuid=None):
+        self.name = name
+        if (text is None and states is None) or (text is not None and states is not None):
+            raise RuntimeError
+        if text is not None:
+            states = (False, text), (True, '\u2714' + text)
+        if default is None:
+            default = states[0][0]
+        self.default = default
+        self.states = OrderedDict(states)
+
+        super(ToggleButton, self).__init__(callback=self._callback, uuid=uuid)
+
+    def post_init(self):
+        self.parent_menu.root_menu.default_data[self.name] = self.default
+
+    def _callback(self, bot, update):
+        data = self.parent_menu.root_menu.data[id_from_update(update)]
+
+        current = data.get(self.name, self.default)
+        keys = list(self.states.keys())
+        index = keys.index(current) + 1
+        if index > len(keys) - 1:
+            index = 0
+        data[self.name] = keys[index]
+
+        return self.parent_menu.callback(bot, update, add_to_stack=False)
+
+    def text(self, update):
+        data = self.parent_menu.root_menu.data[id_from_update(update)]
+        return self.states[data.get(self.name, self.default)]
+
+
+class RadioButton(Button):
+    def __init__(self, name, value, text, enabled=False, uuid=None):
+        self.value = value
+        self.name = name
+
+        super(RadioButton, self).__init__(callback=self._callback, uuid=uuid)
+
+        if isinstance(text, str_type):
+            self._text = ('\u26aa' + text, '\U0001f518' + text)
+        else:
+            self._text = text
+
+        self.enabled = enabled
+
+    def post_init(self):
+        default_data = self.parent_menu.root_menu.default_data
+        if self.name in default_data:
+            if default_data[self.name] is not None:
+                return
+        if self.enabled is True:
+            default_data[self.name] = self.value
+        else:
+            default_data[self.name] = None
+
+    def _callback(self, bot, update):
+        data = self.parent_menu.root_menu.data[id_from_update(update)]
+        data[self.name] = self.value
+
+        return self.parent_menu.callback(bot, update, add_to_stack=False)
+
+    def text(self, update):
+        data = self.parent_menu.root_menu.data[id_from_update(update)]
+        value = data.get(self.name)
+        if (value is None and self.enabled) or value == self.value:
+            return self._text[1]  # True
+        return self._text[0]  # False
 
 
 class MenuHandler(Handler):
     def __init__(self, menu):
         self.menu = menu
         self.buttons = {}
-        self.collect_buttons(self.menu)
 
+        menu.data = defaultdict(dict)
+        menu.default_data = dict()
         menu.stack = defaultdict(list)
 
+        self.collect_buttons(self.menu)
+
         super(MenuHandler, self).__init__(
             pass_update_queue=None,
             pass_job_queue=None,
@@ -177,9 +267,10 @@ def collect_buttons(self, menu):
         menu.root_menu = self.menu
         for button in chain.from_iterable(menu.get_buttons()):
             button.parent_menu = menu
-            if button.name not in self.buttons and (button.callback is not None or
+            button.post_init()
+            if button.uuid not in self.buttons and (button.callback is not None or
                                                     button.menu is not None):
-                self.buttons[button.name] = button
+                self.buttons[button.uuid] = button
                 if button.menu is not None:
                     self.collect_buttons(button.menu)
 
