Added utils.helpers.effective_message_type (python-telegram-bot#826)

JosXa · tsnoam · commit 9338dc4697b9 · 2018-02-15T11:21:19.000+02:00
diff --git a/telegram/message.py b/telegram/message.py
@@ -195,6 +195,14 @@ class Message(TelegramObject):
 
     _effective_attachment = _UNDEFINED
 
+    ATTACHMENT_TYPES = ['audio', 'game', 'document', 'photo', 'sticker', 'video', 'voice',
+                        'video_note', 'contact', 'location', 'venue', 'invoice',
+                        'successful_payment']
+    MESSAGE_TYPES = ['text', 'new_chat_members', 'new_chat_title', 'new_chat_photo',
+                     'delete_chat_photo', 'group_chat_created', 'supergroup_chat_created',
+                     'channel_chat_created', 'migrate_to_chat_id', 'migrate_from_chat_id',
+                     'pinned_message'] + ATTACHMENT_TYPES
+
     def __init__(self,
                  message_id,
                  from_user,
@@ -354,11 +362,9 @@ def effective_attachment(self):
         if self._effective_attachment is not _UNDEFINED:
             return self._effective_attachment
 
-        for i in (self.audio, self.game, self.document, self.photo, self.sticker,
-                  self.video, self.voice, self.video_note, self.contact, self.location,
-                  self.venue, self.invoice, self.successful_payment):
-            if i:
-                self._effective_attachment = i
+        for i in Message.ATTACHMENT_TYPES:
+            if getattr(self, i, None):
+                self._effective_attachment = getattr(self, i)
                 break
         else:
             self._effective_attachment = None
diff --git a/telegram/utils/helpers.py b/telegram/utils/helpers.py
@@ -113,3 +113,34 @@ def mention_markdown(user_id, name):
     """
     if isinstance(user_id, int):
         return '[{}](tg://user?id={})'.format(escape_markdown(name), user_id)
+
+
+def effective_message_type(entity):
+    """
+    Extracts the type of message as a string identifier from a :class:`telegram.Message` or a
+    :class:`telegram.Update`.
+
+    Args:
+        entity (:obj:`Update` | :obj:`Message`) The ``update`` or ``message`` to extract from
+
+    Returns:
+        str: One of ``Message.MESSAGE_TYPES``
+
+    """
+
+    # Importing on file-level yields cyclic Import Errors
+    from telegram import Message
+    from telegram import Update
+
+    if isinstance(entity, Message):
+        message = entity
+    elif isinstance(entity, Update):
+        message = entity.effective_message
+    else:
+        raise TypeError("entity is not Message or Update (got: {})".format(type(entity)))
+
+    for i in Message.MESSAGE_TYPES:
+        if getattr(message, i, None):
+            return i
+
+    return None
diff --git a/tests/test_helpers.py b/tests/test_helpers.py
@@ -16,7 +16,11 @@
 #
 # You should have received a copy of the GNU Lesser Public License
 # along with this program.  If not, see [http://www.gnu.org/licenses/].
+from telegram import Update
 
+from telegram import Sticker
+from telegram import User
+from telegram.message import Message
 from telegram.utils import helpers
 
 
@@ -26,3 +30,28 @@ def test_escape_markdown(self):
         expected_str = '\*bold\*, \_italic\_, \`code\`, \[text\_link](http://github.com/)'
 
         assert expected_str == helpers.escape_markdown(test_str)
+
+    def test_effective_message_type(self):
+        test_message = Message(message_id=1,
+                               from_user=None,
+                               date=None,
+                               chat=None)
+
+        test_message.text = 'Test'
+        assert helpers.effective_message_type(test_message) == 'text'
+        test_message.text = None
+
+        test_message.sticker = Sticker('sticker_id', 50, 50)
+        assert helpers.effective_message_type(test_message) == 'sticker'
+        test_message.sticker = None
+
+        test_message.new_chat_members = [User(55, 'new_user', False)]
+        assert helpers.effective_message_type(test_message) == 'new_chat_members'
+
+        test_update = Update(1)
+        test_message.text = 'Test'
+        test_update.message = test_message
+        assert helpers.effective_message_type(test_update) == 'text'
+
+        empty_update = Update(2)
+        assert helpers.effective_message_type(empty_update) is None
