Merge branch 'feature/edited-messages'

Author: Pietro Albini

botogram/bot.py

@@ -21,6 +21,7 @@
 from . import frozenbot
 from . import shared
 from . import tasks
+from . import messages
 
 
 class Bot(frozenbot.FrozenBot):
@@ -63,6 +64,12 @@ def __init__(self, api_connection):
         # Setup the scheduler
         self._scheduler = tasks.Scheduler()
 
+        # Initialize the list of update processors
+        self._update_processors = {}
+        self.register_update_processor("message", messages.process_message)
+        self.register_update_processor("edited_message",
+                                       messages.process_edited_message)
+
         self._bot_id = str(uuid.uuid4())
 
         self.use(defaults.DefaultComponent())
@@ -130,6 +137,11 @@ def __(func):
             return func
         return __
 
+    def message_edited(self, func):
+        """Add a message edited hook"""
+        self._main_component.add_message_edited_hook(func)
+        return func
+
     def command(self, name, hidden=False):
         """Register a new command"""
         def __(func):
@@ -191,6 +203,14 @@ def run(self, workers=2):
         inst = runner.BotogramRunner(self, workers=workers)
         inst.run()
 
+    def register_update_processor(self, kind, processor):
+        """Register a new update processor"""
+        if kind in self._update_processors:
+            raise NameError("An update processor for \"%s\" updates is "
+                            "already registered" % kind)
+
+        self._update_processors[kind] = processor
+
     def freeze(self):
         """Return a frozen instance of the bot"""
         chains = components.merge_chains(self._main_component,
@@ -202,7 +222,8 @@ def freeze(self):
                                    self.lang, self.itself, self._commands_re,
                                    self._commands, chains, self._scheduler,
                                    self._main_component._component_id,
-                                   self._bot_id, self._shared_memory)
+                                   self._bot_id, self._shared_memory,
+                                   self._update_processors)
 
     @property
     def lang(self):

botogram/components.py

@@ -31,6 +31,7 @@ def __new__(cls, *args, **kwargs):
         self.__memory_preparers = []
         self.__timers = []
         self.__chat_unavailable_hooks = []
+        self.__messages_edited_hooks = []
 
         self._component_id = str(uuid.uuid4())
 
@@ -148,6 +149,14 @@ def add_chat_unavailable_hook(self, func):
         hook = hooks.ChatUnavailableHook(func, self)
         self.__chat_unavailable_hooks.append(hook)
 
+    def add_message_edited_hook(self, func):
+        """Add a new edited message hook"""
+        if not callable(func):
+            raise ValueError("A message edited hook must be callable")
+
+        hook = hooks.MessageEditedHook(func, self)
+        self.__messages_edited_hooks.append(hook)
+
     def _add_no_commands_hook(self, func):
         """Register an hook which will be executed when no commands matches"""
         if not callable(func):
@@ -170,6 +179,7 @@ def _get_chains(self):
             "memory_preparers": [self.__memory_preparers],
             "tasks": [self.__timers],
             "chat_unavalable_hooks": [self.__chat_unavailable_hooks],
+            "messages_edited": [self.__messages_edited_hooks],
         }
 
     def _get_commands(self):

botogram/frozenbot.py

@@ -23,7 +23,7 @@ class FrozenBot:
     def __init__(self, api, about, owner, hide_commands, before_help,
                  after_help, process_backlog, lang, itself, commands_re,
                  commands, chains, scheduler, main_component_id, bot_id,
-                 shared_memory):
+                 shared_memory, update_processors):
         # This attribute should be added with the default setattr, because is
         # needed by the custom setattr
         object.__setattr__(self, "_frozen", False)
@@ -43,6 +43,7 @@ def __init__(self, api, about, owner, hide_commands, before_help,
         self._shared_memory = shared_memory
         self._scheduler = scheduler
         self._chains = chains
+        self._update_processors = update_processors
         self._commands = {name: command.for_bot(self)
                           for name, command in commands.items()}
 
@@ -65,7 +66,7 @@ def __reduce__(self):
             self.before_help, self.after_help, self.process_backlog,
             self.lang, self.itself, self._commands_re, self._commands,
             self._chains, self._scheduler, self._main_component_id,
-            self._bot_id, self._shared_memory,
+            self._bot_id, self._shared_memory, self._update_processors,
         )
         return restore, args
 
@@ -171,21 +172,13 @@ def process(self, update):
         update.set_api(self.api)  # Be sure to use the correct API object
 
         try:
-            for hook in self._chains["messages"]:
-                # Get the correct name of the hook
-                name = hook.name
-                self.logger.debug("Processing update #%s with the %s hook..." %
-                                  (update.update_id, name))
-
-                result = hook.call(self, update)
-                if result is True:
-                    self.logger.debug("Update #%s was just processed by the "
-                                      "%s hook." % (update.update_id, name))
-                    return True
-
-            self.logger.debug("No hook actually processed the #%s update." %
-                              update.update_id)
+            for kind, processor in self._update_processors.items():
+                # Call the processor of the right kind
+                if getattr(update, kind) is None:
+                    continue
 
+                processor(self, self._chains, update)
+                break
         except api_module.ChatUnavailableError as e:
             # Do some sane logging
             self.logger.warning("Chat %s is not available to your bot:" %
@@ -199,8 +192,6 @@ def process(self, update):
                                   e.chat_id))
                 hook.call(self, e.chat_id, e.reason)
 
-        return False
-
     def scheduled_tasks(self, current_time=None, wrap=True):
         """Return a list of tasks scheduled for now"""
         # This provides a better API for the users of the method
@@ -215,6 +206,10 @@ def process():
             return [wrapper(job) for job in tasks]
         return list(tasks)
 
+    def register_update_processor(self, kind, processor):
+        """Register a new update processor"""
+        raise FrozenBotError("Can't register new update processors at runtime")
+
     # This helper manages the translation
 
     def _(self, message, **args):

botogram/hooks.py

@@ -200,6 +200,15 @@ def call(self, bot, chat_id, reason):
                          reason=reason)
 
 
+class MessageEditedHook(Hook):
+    """Underlying hook for @bot.message_edited"""
+
+    def _call(self, bot, update):
+        message = update.edited_message
+        return bot._call(self.func, self.component_id, chat=message.chat,
+                         message=message)
+
+
 class TimerHook(Hook):
     """Underlying hook for a timer"""
 

botogram/messages.py

@@ -0,0 +1,39 @@
+"""
+    botogram.messages
+    Logic for processing messages sent to your bot
+
+    Copyright (c) 2016 Pietro Albini <[email protected]>
+    Released under the MIT license
+"""
+
+
+def process_message(bot, chains, update):
+    """Process a message sent to the bot"""
+    for hook in chains["messages"]:
+        bot.logger.debug("Processing 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)
+
+
+def process_edited_message(bot, chains, update):
+    """Process an edited message"""
+    for hook in chains["messages_edited"]:
+        bot.logger.debug("Processing edited message 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/messages.py

@@ -319,6 +319,7 @@ def from_(self):
         "migrate_to_chat_id": int,
         "migrate_from_chat_id": int,
         "pinned_message": _itself,
+        "edit_date": int,
     }
     replace_keys = {
         "from": "sender",

botogram/objects/updates.py

@@ -21,7 +21,8 @@ class Update(BaseObject):
         "update_id": int,
     }
     optional = {
-        "message": Message
+        "message": Message,
+        "edited_message": Message,
     }
     _check_equality_ = "update_id"
 

docs/api/bot.rst

@@ -163,6 +163,31 @@ components.
       :param bool multiple: If the function should be called multiple times on
          multiple matches.
 
+   .. py:decoratormethod:: message_edited
+
+      All the functions decorated with this method will be called when an user
+      edits a message the bot knows about. This allows you, for example, to
+      update the preview of a message if the user edits the request, or to
+      enforce a no-edits policy on groups by banning whoever edits a message.
+
+      You can :ref:`request the following arguments <bot-structure-hooks-args>`
+      in the decorated functions:
+
+      * **chat**: the chat in which the message was originally sent (instance
+        of :py:class:`~botogram.Chat`)
+
+      * **message**: the edited message (instance of
+        :py:class:`~botogram.Message`)
+
+      .. code-block:: python
+
+         @bot.message_edited
+         def no_edits(chat, message):
+             message.reply("You can't edit messages! Bye.")
+             chat.ban(message.sender)
+
+      .. versionadded:: 0.3
+
    .. py:decoratormethod:: command(name, [hidden=False])
 
       This decorator register a new command, and calls the decorated function

docs/api/components.rst

@@ -127,6 +127,37 @@ about how to create them in the ":ref:`custom-components`" chapter.
       :param bool multiple: If the function should be called multiple times on
          multiple matches.
 
+   .. py:method:: add_message_edited_hook(func)
+
+      All the functions provided to this method will be called when an user
+      edits a message the bot knows about. This allows you, for example, to
+      update the preview of a message if the user edits the request, or to
+      enforce a no-edits policy on groups by banning whoever edits a message.
+
+      You can :ref:`request the following arguments <bot-structure-hooks-args>`
+      in the provided functions:
+
+      * **chat**: the chat in which the message was orignally sent (instance of
+        :py:class:`~botogram.Chat`)
+      * **message**: the edited message (instance of
+        :py:class:`~botogram.Message`)
+
+      .. code-block:: python
+
+         class NoEditsComponent(botogram.Component):
+             component_name = "noedits"
+
+             def __init__(self):
+                 self.add_message_edited_hook(self.no_edits)
+
+             def no_edits(self, chat, message):
+                 message.reply("You can't edit messages! Bye.")
+                 chat.ban(message.sender)
+
+      :param callable func: The function you want to use.
+
+      .. versionadded:: 0.3
+
    .. py:method:: add_command(name, func, [hidden=False])
 
       This function registers a new command, and calls the provided function

docs/changelog.rst

@@ -45,6 +45,11 @@ New features
    * New argument ``hidden`` on :py:meth:`botogram.Bot.command`
    * New argument ``hidden`` on :py:meth:`botogram.Component.add_command`
 
+* Add support for processing edited messages:
+
+   * New decorator :py:meth:`botogram.Bot.message_edited`
+   * New method :py:meth:`botogram.Component.add_message_edited_hook`
+
 * Added new attribute :py:attr:`botogram.Message.pinned_message`
 * Added new attribute :py:attr:`botogram.Sticker.emoji`
 * Added new attribute :py:attr:`botogram.Chat.admins`