Add @bot.process_poll hook

Author: MarcoBuster

botogram/bot.py

@@ -92,6 +92,7 @@ def __init__(self, api_connection):
         self.register_update_processor("edited_channel_post",
                                        messages.process_channel_post_edited)
         self.register_update_processor("callback_query", callbacks.process)
+        self.register_update_processor("poll", messages.process_poll_update)
 
         self._bot_id = str(uuid.uuid4())
 
@@ -142,6 +143,11 @@ def process_message(self, func):
         self._main_component.add_process_message_hook(func)
         return func
 
+    def poll_update(self, func):
+        """Add a poll update hook"""
+        self._main_component.add_poll_update_hook(func)
+        return func
+
     def message_equals(self, string, ignore_case=True):
         """Add a message equals hook"""
         def __(func):

botogram/components.py

@@ -47,6 +47,7 @@ def __new__(cls, *args, **kwargs):
         self.__messages_edited_hooks = []
         self.__channel_post_hooks = []
         self.__channel_post_edited_hooks = []
+        self.__poll_update_hooks = []
 
         self._component_id = str(uuid.uuid4())
 
@@ -75,6 +76,14 @@ def add_process_message_hook(self, func):
         hook = hooks.ProcessMessageHook(func, self)
         self.__processors.append(hook)
 
+    def add_poll_update_hook(self, func):
+        """Add a poll update hook"""
+        if not callable(func):
+            raise ValueError("A poll update hook must be callable")
+
+        hook = hooks.PollUpdateHook(func, self)
+        self.__poll_update_hooks.append(hook)
+
     def add_message_equals_hook(self, string, func, ignore_case=True):
         """Add a message equals hook"""
         if not callable(func):
@@ -220,6 +229,7 @@ def _get_chains(self):
         ]
         return {
             "messages": messages,
+            "poll_updates": [self.__poll_update_hooks],
             "memory_preparers": [self.__memory_preparers],
             "tasks": [self.__timers],
             "chat_unavalable_hooks": [self.__chat_unavailable_hooks],

botogram/frozenbot.py

@@ -111,6 +111,10 @@ def process_message(self, func):
         """Register a message processor hook"""
         raise FrozenBotError("Can't add hooks to a bot at runtime")
 
+    def poll_update(self, func):
+        """Register a poll update hook"""
+        raise FrozenBotError("Can't add hooks to a bot at runtime")
+
     def message_equals(self, string, ignore_case=True):
         """Add a message equals hook"""
         raise FrozenBotError("Can't add hooks to a bot at runtime")

botogram/hooks.py

@@ -81,6 +81,13 @@ class ProcessMessageHook(Hook):
     pass
 
 
+class PollUpdateHook(Hook):
+    """Underlying hook for @bot.poll_update"""
+
+    def _call(self, bot, update):
+        return bot._call(self.func, self.component_id, poll=update.poll)
+
+
 class MemoryPreparerHook(Hook):
     """Underlying hook for @bot.prepare_memory"""
 

botogram/messages.py

@@ -81,3 +81,19 @@ def process_channel_post_edited(bot, chains, update):
 
     bot.logger.debug("No hook actually processed the #%s update." %
                      update.update_id)
+
+
+def process_poll_update(bot, chains, update):
+    """Process a poll update"""
+    for hook in chains["poll_updates"]:
+        bot.logger.debug("Processing poll update in update #%s with"
+                         "the hook %s..." % (update.update_id, hook.name))
+
+        result = hook.call(bot, update)
+        if result is True:
+            bot.logger.debug("Update %s was just processed by the %s hook." %
+                             (update.update_id, hook.name))
+            return
+
+    bot.logger.debug("No hook actually processed the #%s update." %
+                     update.update_id)

botogram/objects/updates.py

@@ -22,6 +22,7 @@
 
 from .callbacks import CallbackQuery
 from .messages import Message
+from .polls import Poll
 
 
 class Update(BaseObject):
@@ -41,6 +42,7 @@ class Update(BaseObject):
         "channel_post": Message,
         "edited_channel_post": Message,
         "callback_query": CallbackQuery,
+        "poll": Poll,
     }
     _check_equality_ = "update_id"