Merge pull request python-telegram-bot#409 from python-telegram-bot/entities-filter

jsmnbom · web-flow · commit 5f6138b06be0 · 2016-09-24T14:02:01.000+02:00
Add entities filter
diff --git a/telegram/ext/messagehandler.py b/telegram/ext/messagehandler.py
@@ -85,6 +85,23 @@ def status_update(message):
     def forwarded(message):
         return bool(message.forward_date)
 
+    @staticmethod
+    def entity(entity_type):
+        """Filters messages to only allow those which have a :class:`telegram.MessageEntity`
+        where their `type` matches `entity_type`.
+
+        Args:
+            entity_type: Entity type to check for. All types can be found as constants
+                in :class:`telegram.MessageEntity`.
+
+        Returns: function to use as filter
+        """
+
+        def entities_filter(message):
+            return any([entity.type == entity_type for entity in message.entities])
+
+        return entities_filter
+
 
 class MessageHandler(Handler):
     """
@@ -141,7 +158,8 @@ def handle_update(self, update, dispatcher):
 
         return self.callback(dispatcher.bot, update, **optional_args)
 
-    # old non-PEP8 Handler methods
+# old non-PEP8 Handler methods
+
     m = "telegram.MessageHandler."
     checkUpdate = deprecate(check_update, m + "checkUpdate", m + "check_update")
     handleUpdate = deprecate(handle_update, m + "handleUpdate", m + "handle_update")
diff --git a/tests/test_filters.py b/tests/test_filters.py
@@ -23,10 +23,11 @@
 import sys
 import unittest
 from datetime import datetime
+import functools
 
 sys.path.append('.')
 
-from telegram import Message, User, Chat
+from telegram import Message, User, Chat, MessageEntity
 from telegram.ext import Filters
 from tests.base import BaseTest
 
@@ -150,6 +151,21 @@ def test_filters_status_update(self):
         self.assertTrue(Filters.status_update(self.message))
         self.message.pinned_message = None
 
+    def test_entities_filter(self):
+        e = functools.partial(MessageEntity, offset=0, length=0)
+
+        self.message.entities = [e(MessageEntity.MENTION)]
+        self.assertTrue(Filters.entity(MessageEntity.MENTION)(self.message))
+
+        self.message.entities = []
+        self.assertFalse(Filters.entity(MessageEntity.MENTION)(self.message))
+
+        self.message.entities = [e(MessageEntity.BOLD)]
+        self.assertFalse(Filters.entity(MessageEntity.MENTION)(self.message))
+
+        self.message.entities = [e(MessageEntity.BOLD), e(MessageEntity.MENTION)]
+        self.assertTrue(Filters.entity(MessageEntity.MENTION)(self.message))
+
 
 if __name__ == '__main__':
     unittest.main()
