Fix HTML and Markdown unparsing

delivrance · delivrance · commit 8a2416665b5c · 2022-04-25T08:30:26.000+02:00
diff --git a/pyrogram/parser/html.py b/pyrogram/parser/html.py
@@ -24,6 +24,7 @@
 
 import pyrogram
 from pyrogram import raw
+from pyrogram.enums import MessageEntityType
 from pyrogram.errors import PeerIdInvalid
 from . import utils
 
@@ -155,17 +156,21 @@ def unparse(text: str, entities: list):
             start = entity.offset
             end = start + entity.length
 
-            if entity_type in ("bold", "italic", "underline", "strikethrough"):
-                start_tag = f"<{entity_type[0]}>"
-                end_tag = f"</{entity_type[0]}>"
-            elif entity_type in ("code", "pre", "blockquote", "spoiler"):
-                start_tag = f"<{entity_type}>"
-                end_tag = f"</{entity_type}>"
-            elif entity_type == "text_link":
+            if entity_type in (MessageEntityType.BOLD, MessageEntityType.ITALIC, MessageEntityType.UNDERLINE,
+                               MessageEntityType.STRIKETHROUGH):
+                name = entity_type.name[0].lower()
+                start_tag = f"<{name}>"
+                end_tag = f"</{name}>"
+            elif entity_type in (MessageEntityType.CODE, MessageEntityType.PRE, MessageEntityType.BLOCKQUOTE,
+                                 MessageEntityType.SPOILER):
+                name = entity_type.name.lower()
+                start_tag = f"<{name}>"
+                end_tag = f"</{name}>"
+            elif entity_type == MessageEntityType.TEXT_LINK:
                 url = entity.url
                 start_tag = f'<a href="{url}">'
                 end_tag = "</a>"
-            elif entity_type == "text_mention":
+            elif entity_type == MessageEntityType.TEXT_MENTION:
                 user = entity.user
                 start_tag = f'<a href="tg://user?id={user.id}">'
                 end_tag = "</a>"
diff --git a/pyrogram/parser/markdown.py b/pyrogram/parser/markdown.py
@@ -21,6 +21,7 @@
 from typing import Optional
 
 import pyrogram
+from pyrogram.enums import MessageEntityType
 from . import utils
 from .html import HTML
 
@@ -119,25 +120,25 @@ def unparse(text: str, entities: list):
             start = entity.offset
             end = start + entity.length
 
-            if entity_type == "bold":
+            if entity_type == MessageEntityType.BOLD:
                 start_tag = end_tag = BOLD_DELIM
-            elif entity_type == "italic":
+            elif entity_type == MessageEntityType.ITALIC:
                 start_tag = end_tag = ITALIC_DELIM
-            elif entity_type == "underline":
+            elif entity_type == MessageEntityType.UNDERLINE:
                 start_tag = end_tag = UNDERLINE_DELIM
-            elif entity_type == "strikethrough":
+            elif entity_type == MessageEntityType.STRIKETHROUGH:
                 start_tag = end_tag = STRIKE_DELIM
-            elif entity_type == "code":
+            elif entity_type == MessageEntityType.CODE:
                 start_tag = end_tag = CODE_DELIM
-            elif entity_type in ("pre", "blockquote"):
+            elif entity_type in (MessageEntityType.PRE, MessageEntityType.BLOCKQUOTE):
                 start_tag = end_tag = PRE_DELIM
-            elif entity_type == "spoiler":
+            elif entity_type == MessageEntityType.SPOILER:
                 start_tag = end_tag = SPOILER_DELIM
-            elif entity_type == "text_link":
+            elif entity_type == MessageEntityType.TEXT_LINK:
                 url = entity.url
                 start_tag = "["
                 end_tag = f"]({url})"
-            elif entity_type == "text_mention":
+            elif entity_type == MessageEntityType.TEXT_MENTION:
                 user = entity.user
                 start_tag = "["
                 end_tag = f"](tg://user?id={user.id})"
