diff --git a/userge/plugins/misc/download.py b/userge/plugins/misc/download.py index b2136f841..df58dd193 100644 --- a/userge/plugins/misc/download.py +++ b/userge/plugins/misc/download.py @@ -20,7 +20,7 @@ from pySmartDL import SmartDL from userge import userge, Message, Config -from userge.utils import progress, humanbytes, extract_urls +from userge.utils import progress, humanbytes, extract_entities from userge.utils.exceptions import ProcessCanceled LOGGER = userge.getLogger(__name__) @@ -151,7 +151,7 @@ async def tg_download( """ download from tg file """ if not to_download.media: dl_loc, mite = [], 0 - ets = extract_urls(to_download) + ets = extract_entities(to_download, ["url", "text_link"]) if len(ets) == 0: raise Exception("nothing found to download") for uarl in ets: diff --git a/userge/utils/__init__.py b/userge/utils/__init__.py index fe4051a21..d389fa6e8 100644 --- a/userge/utils/__init__.py +++ b/userge/utils/__init__.py @@ -22,4 +22,4 @@ take_screen_shot, parse_buttons, is_command, - extract_urls) + extract_entities) diff --git a/userge/utils/tools.py b/userge/utils/tools.py index 61b0d0b34..bf932e290 100644 --- a/userge/utils/tools.py +++ b/userge/utils/tools.py @@ -222,18 +222,42 @@ def is_command(cmd: str) -> bool: return is_cmd -def extract_urls(message: Message) -> List[str]: +def extract_entities(message: Message, typeofentity: List[str]) -> List[str]: + """ gets a message and returns a list of entity_type in the message + """ tero = [] entities = message.entities or message.caption_entities or [] text = message.text or message.caption or "" for entity in entities: url = None - if entity.type == "text_link": - url = entity.url - elif entity.type == "url": + cet = entity.type + if entity.type in [ + "url", + "mention", + "hashtag", + "cashtag", + "bot_command", + "url", + "email", + "phone_number", + "bold", + "italic", + "underline", + "strikethrough", + "spoiler", + "code", + "pre", + ]: offset = entity.offset length = entity.length url = text[offset:offset + length] - if url: + + elif entity.type == "text_link": + url = entity.url + + elif entity.type == "text_mention": + url = entity.user + + if url and cet in typeofentity: tero.append(url) return tero