Added utils.helpers.effective_message_type (#826)

This commit is contained in:
Joscha Götzer 2018-02-15 10:21:19 +01:00 committed by Noam Meltzer
parent 063704c0b9
commit 9338dc4697
3 changed files with 71 additions and 5 deletions

View file

@ -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 @@ class Message(TelegramObject):
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

View file

@ -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

View file

@ -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 @@ class TestHelpers(object):
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