diff --git a/telethon/tl/custom/message.py b/telethon/tl/custom/message.py index 93b91318..7f373284 100644 --- a/telethon/tl/custom/message.py +++ b/telethon/tl/custom/message.py @@ -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): diff --git a/telethon/utils.py b/telethon/utils.py index 312b47ac..d8b4b97b 100644 --- a/telethon/utils.py +++ b/telethon/utils.py @@ -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):