Fix get_entities_text and allow filtering by type

This commit is contained in:
Lonami Exo 2018-06-07 10:46:32 +02:00
parent 4bdc28a775
commit 0418f7e375
2 changed files with 27 additions and 15 deletions

View File

@ -495,13 +495,31 @@ class Message:
return self._client.download_media(self.original_message,
*args, **kwargs)
def get_entities_text(self):
def get_entities_text(self, cls=None):
"""
Returns a list of tuples [(:tl:`MessageEntity`, `str`)], the string
being the inner text of the message entity (like bold, italics, etc).
Args:
cls (`type`):
Returns entities matching this type only. For example,
the following will print the text for all ``code`` entities:
>>> from telethon.tl.types import MessageEntityCode
>>>
>>> m = Message(...)
>>> for _, inner_text in m.get_entities_text(MessageEntityCode):
>>> print(inner_text)
"""
texts = get_inner_text(self.original_message.message,
self.original_message.entities)
if cls and self.original_message.entities:
texts = get_inner_text(
self.original_message.message,
[c for c in self.original_message.entities
if isinstance(c, cls)]
)
else:
texts = get_inner_text(self.original_message.message,
self.original_message.entities)
return list(zip(self.original_message.entities, texts))
def click(self, i=None, j=None, *, text=None, filter=None):

View File

@ -527,29 +527,23 @@ def del_surrogate(text):
return text.encode('utf-16', 'surrogatepass').decode('utf-16')
def get_inner_text(text, entity):
def get_inner_text(text, entities):
"""
Gets the inner text that's surrounded by the given entity or entities.
Gets the inner text that's surrounded by the given entities.
For instance: text = 'hey!', entity = MessageEntityBold(2, 2) -> 'y!'.
:param text: the original text.
:param entity: the entity or entities that must be matched.
:param text: the original text.
:param entities: the entity or entities that must be matched.
:return: a single result or a list of the text surrounded by the entities.
"""
if isinstance(entity, TLObject):
entity = (entity,)
multiple = True
else:
multiple = False
text = add_surrogate(text)
result = []
for e in entity:
for e in entities:
start = e.offset
end = e.offset + e.length
result.append(del_surrogate(text[start:end]))
return result if multiple else result[0]
return result
def get_peer_id(peer):