Add type hint in parser

pyrogram · mo1ein · Jun 22, 2023 · Jun 22, 2023 · Jun 22, 2023 · Jun 22, 2023
commit 79ca38c37c228bf70df29da5820b01687d4159c9
diff --git a/pyrogram/parser/html.py b/pyrogram/parser/html.py
@@ -20,7 +20,7 @@
 import logging
 import re
 from html.parser import HTMLParser
-from typing import Optional
+from typing import Optional, Union
 
 import pyrogram
 from pyrogram import raw
@@ -34,7 +34,7 @@
 class Parser(HTMLParser):
     MENTION_RE = re.compile(r"tg://user\?id=(\d+)")
 
-    def __init__(self, client: "pyrogram.Client"):
+    def __init__(self, client: "pyrogram.Client") -> None:
         super().__init__()
 
         self.client = client
@@ -43,7 +43,7 @@ def __init__(self, client: "pyrogram.Client"):
         self.entities = []
         self.tag_entities = {}
 
-    def handle_starttag(self, tag, attrs):
+    def handle_starttag(self, tag: str, attrs: list) -> None:
         attrs = dict(attrs)
         extra = {}
 
@@ -87,7 +87,7 @@ def handle_starttag(self, tag, attrs):
 
         self.tag_entities[tag].append(entity(offset=len(self.text), length=0, **extra))
 
-    def handle_data(self, data):
+    def handle_data(self, data: str) -> None:
         data = html.unescape(data)
 
         for entities in self.tag_entities.values():
@@ -96,7 +96,7 @@ def handle_data(self, data):
 
         self.text += data
 
-    def handle_endtag(self, tag):
+    def handle_endtag(self, tag: str) -> None:
         try:
             self.entities.append(self.tag_entities[tag].pop())
         except (KeyError, IndexError):
@@ -113,10 +113,10 @@ def error(self, message):
 
 
 class HTML:
-    def __init__(self, client: Optional["pyrogram.Client"]):
+    def __init__(self, client: Optional["pyrogram.Client"]) -> None:
         self.client = client
 
-    async def parse(self, text: str):
+    async def parse(self, text: str) -> dict:
         # Strip whitespaces from the beginning and the end, but preserve closing tags
         text = re.sub(r"^\s*(<[\w<>=\s\"]*>)\s*", r"\1", text)
         text = re.sub(r"\s*(</[\w</>]*>)\s*$", r"\1", text)
@@ -154,8 +154,8 @@ async def parse(self, text: str):
         }
 
     @staticmethod
-    def unparse(text: str, entities: list):
-        def parse_one(entity):
+    def unparse(text: str, entities: list) -> str:
+        def parse_one(entity) -> Union[tuple, None]:
             """
             Parses a single entity and returns (start_tag, start), (end_tag, end)
             """

diff --git a/pyrogram/parser/markdown.py b/pyrogram/parser/markdown.py
@@ -56,10 +56,10 @@
 
 
 class Markdown:
-    def __init__(self, client: Optional["pyrogram.Client"]):
+    def __init__(self, client: Optional["pyrogram.Client"]) -> None:
         self.html = HTML(client)
 
-    async def parse(self, text: str, strict: bool = False):
+    async def parse(self, text: str, strict: bool = False) -> dict:
         if strict:
             text = html.escape(text)
 
@@ -116,7 +116,7 @@ async def parse(self, text: str, strict: bool = False):
         return await self.html.parse(text)
 
     @staticmethod
-    def unparse(text: str, entities: list):
+    def unparse(text: str, entities: list) -> str:
         text = utils.add_surrogates(text)
 
         entities_offsets = []

diff --git a/pyrogram/parser/parser.py b/pyrogram/parser/parser.py
@@ -16,7 +16,7 @@
 #  You should have received a copy of the GNU Lesser General Public License
 #  along with Pyrogram.  If not, see <http://www.gnu.org/licenses/>.
 
-from typing import Optional
+from typing import Optional, Union
 
 import pyrogram
 from pyrogram import enums
@@ -30,7 +30,7 @@ def __init__(self, client: Optional["pyrogram.Client"]):
         self.html = HTML(client)
         self.markdown = Markdown(client)
 
-    async def parse(self, text: str, mode: Optional[enums.ParseMode] = None):
+    async def parse(self, text: str, mode: Optional[enums.ParseMode] = None) -> Union[str, dict]:
         text = str(text if text else "").strip()
 
         if mode is None:
@@ -54,8 +54,7 @@ async def parse(self, text: str, mode: Optional[enums.ParseMode] = None):
         raise ValueError(f'Invalid parse mode "{mode}"')
 
     @staticmethod
-    def unparse(text: str, entities: list, is_html: bool):
+    def unparse(text: str, entities: list, is_html: bool) -> str:
         if is_html:
             return HTML.unparse(text, entities)
-        else:
-            return Markdown.unparse(text, entities)
+        return Markdown.unparse(text, entities)
diff --git a/pyrogram/parser/utils.py b/pyrogram/parser/utils.py
@@ -32,10 +32,10 @@ def add_surrogates(text):
     )
 
 
-def remove_surrogates(text):
+def remove_surrogates(text: str) -> str:
     # Replace each surrogate pair with a SMP code point
     return text.encode("utf-16", "surrogatepass").decode("utf-16")
 
 
-def replace_once(source: str, old: str, new: str, start: int):
+def replace_once(source: str, old: str, new: str, start: int) -> str:
     return source[:start] + source[start:].replace(old, new, 1)