improve extract_urls function, and rename the function (UsergeTeam#462)

JustANothing · Mar 5, 2022 · 19dde91 · 19dde91
1 parent d6bb4a3
commit 19dde91
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 8 deletions.
diff --git 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
@@ -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
@@ -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