Skip to content

Commit 57ee4e5

Browse files
committed
Add @bot.process_poll hook
1 parent 643cb9b commit 57ee4e5

File tree

6 files changed

+45
-0
lines changed

6 files changed

+45
-0
lines changed

botogram/bot.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -92,6 +92,7 @@ def __init__(self, api_connection):
9292
self.register_update_processor("edited_channel_post",
9393
messages.process_channel_post_edited)
9494
self.register_update_processor("callback_query", callbacks.process)
95+
self.register_update_processor("poll", messages.process_poll_update)
9596

9697
self._bot_id = str(uuid.uuid4())
9798

@@ -142,6 +143,11 @@ def process_message(self, func):
142143
self._main_component.add_process_message_hook(func)
143144
return func
144145

146+
def poll_update(self, func):
147+
"""Add a poll update hook"""
148+
self._main_component.add_poll_update_hook(func)
149+
return func
150+
145151
def message_equals(self, string, ignore_case=True):
146152
"""Add a message equals hook"""
147153
def __(func):

botogram/components.py

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,7 @@ def __new__(cls, *args, **kwargs):
4747
self.__messages_edited_hooks = []
4848
self.__channel_post_hooks = []
4949
self.__channel_post_edited_hooks = []
50+
self.__poll_update_hooks = []
5051

5152
self._component_id = str(uuid.uuid4())
5253

@@ -75,6 +76,14 @@ def add_process_message_hook(self, func):
7576
hook = hooks.ProcessMessageHook(func, self)
7677
self.__processors.append(hook)
7778

79+
def add_poll_update_hook(self, func):
80+
"""Add a poll update hook"""
81+
if not callable(func):
82+
raise ValueError("A poll update hook must be callable")
83+
84+
hook = hooks.PollUpdateHook(func, self)
85+
self.__poll_update_hooks.append(hook)
86+
7887
def add_message_equals_hook(self, string, func, ignore_case=True):
7988
"""Add a message equals hook"""
8089
if not callable(func):
@@ -220,6 +229,7 @@ def _get_chains(self):
220229
]
221230
return {
222231
"messages": messages,
232+
"poll_updates": [self.__poll_update_hooks],
223233
"memory_preparers": [self.__memory_preparers],
224234
"tasks": [self.__timers],
225235
"chat_unavalable_hooks": [self.__chat_unavailable_hooks],

botogram/frozenbot.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -111,6 +111,10 @@ def process_message(self, func):
111111
"""Register a message processor hook"""
112112
raise FrozenBotError("Can't add hooks to a bot at runtime")
113113

114+
def poll_update(self, func):
115+
"""Register a poll update hook"""
116+
raise FrozenBotError("Can't add hooks to a bot at runtime")
117+
114118
def message_equals(self, string, ignore_case=True):
115119
"""Add a message equals hook"""
116120
raise FrozenBotError("Can't add hooks to a bot at runtime")

botogram/hooks.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -81,6 +81,13 @@ class ProcessMessageHook(Hook):
8181
pass
8282

8383

84+
class PollUpdateHook(Hook):
85+
"""Underlying hook for @bot.poll_update"""
86+
87+
def _call(self, bot, update):
88+
return bot._call(self.func, self.component_id, poll=update.poll)
89+
90+
8491
class MemoryPreparerHook(Hook):
8592
"""Underlying hook for @bot.prepare_memory"""
8693

botogram/messages.py

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -81,3 +81,19 @@ def process_channel_post_edited(bot, chains, update):
8181

8282
bot.logger.debug("No hook actually processed the #%s update." %
8383
update.update_id)
84+
85+
86+
def process_poll_update(bot, chains, update):
87+
"""Process a poll update"""
88+
for hook in chains["poll_updates"]:
89+
bot.logger.debug("Processing poll update in update #%s with"
90+
"the hook %s..." % (update.update_id, hook.name))
91+
92+
result = hook.call(bot, update)
93+
if result is True:
94+
bot.logger.debug("Update %s was just processed by the %s hook." %
95+
(update.update_id, hook.name))
96+
return
97+
98+
bot.logger.debug("No hook actually processed the #%s update." %
99+
update.update_id)

botogram/objects/updates.py

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

2323
from .callbacks import CallbackQuery
2424
from .messages import Message
25+
from .polls import Poll
2526

2627

2728
class Update(BaseObject):
@@ -41,6 +42,7 @@ class Update(BaseObject):
4142
"channel_post": Message,
4243
"edited_channel_post": Message,
4344
"callback_query": CallbackQuery,
45+
"poll": Poll,
4446
}
4547
_check_equality_ = "update_id"
4648

0 commit comments

Comments
 (0)