Skip to content

Commit

Permalink
Add a caption_entity filter for filtering caption entities (#1068)
Browse files Browse the repository at this point in the history
* Add a caption_entity filter for filtering caption entities

* remove unneeded list comprehensions
  • Loading branch information
PaulSonOfLars authored and jh0ker committed Apr 20, 2018
1 parent 39c679e commit b5196f0
Show file tree
Hide file tree
Showing 2 changed files with 37 additions and 1 deletion.
23 changes: 22 additions & 1 deletion telegram/ext/filters.py
Original file line number Diff line number Diff line change
Expand Up @@ -541,7 +541,28 @@ def __init__(self, entity_type):
self.name = 'Filters.entity({})'.format(self.entity_type)

def filter(self, message):
return any([entity.type == self.entity_type for entity in message.entities])
return any(entity.type == self.entity_type for entity in message.entities)

class caption_entity(BaseFilter):
"""
Filters media messages to only allow those which have a :class:`telegram.MessageEntity`
where their `type` matches `entity_type`.
Examples:
Example ``MessageHandler(Filters.caption_entity("hashtag"), callback_method)``
Args:
entity_type: Caption Entity type to check for. All types can be found as constants
in :class:`telegram.MessageEntity`.
"""

def __init__(self, entity_type):
self.entity_type = entity_type
self.name = 'Filters.caption_entity({})'.format(self.entity_type)

def filter(self, message):
return any(entity.type == self.entity_type for entity in message.caption_entities)

class _Private(BaseFilter):
name = 'Filters.private'
Expand Down
15 changes: 15 additions & 0 deletions tests/test_filters.py
Original file line number Diff line number Diff line change
Expand Up @@ -310,6 +310,21 @@ def test_entities_filter(self, message, message_entity):
second = MessageEntity.de_json(second, None)
message.entities = [message_entity, second]
assert Filters.entity(message_entity.type)(message)
assert not Filters.caption_entity(message_entity.type)(message)

def test_caption_entities_filter(self, message, message_entity):
message.caption_entities = [message_entity]
assert Filters.caption_entity(message_entity.type)(message)

message.caption_entities = []
assert not Filters.caption_entity(MessageEntity.MENTION)(message)

second = message_entity.to_dict()
second['type'] = 'bold'
second = MessageEntity.de_json(second, None)
message.caption_entities = [message_entity, second]
assert Filters.caption_entity(message_entity.type)(message)
assert not Filters.entity(message_entity.type)(message)

def test_private_filter(self, message):
assert Filters.private(message)
Expand Down

0 comments on commit b5196f0

Please sign in to comment.