refine comments and add type hints

Author: duccuong.le

paddleocr_convert/main.py

@@ -84,7 +84,9 @@ def convert_to_onnx(self, unzip_result: object, save_onnx_path: str) -> None:
         """借助 :code:`paddle2onnx` 工具转换模型为onnx格式
 
         Args:
-            model_dir (str): 保存paddle格式模型所在目录
+            unzip_result (object):
+                - model_dir (Path): 解压后模型保存路径。
+                - my_files (set): 解压得到的文件名集合。
             save_onnx_path (str): 保存的onnx全路径
         """
         model_dir = unzip_result.get("model_dir", "")

paddleocr_convert/utils.py

@@ -5,14 +5,19 @@
 import re
 import tarfile
 from pathlib import Path
-from typing import Union
+from typing import TypedDict, Set, Union
 
 import requests
 from tqdm import tqdm
 
 InputType = Union[str, Path]
 
 
+class UnzipResult(TypedDict):
+    model_dir: Path
+    my_files: Set[str]
+
+
 class DownloadModelError(Exception):
     pass
 
@@ -52,7 +57,7 @@ def download_file(url: str, save_dir: InputType) -> Path:
 
     save_path = Path(save_dir) / Path(url).name
     with tqdm(
-        total=total_size_in_bytes, unit="iB", unit_scale=True, desc="Downloading"
+            total=total_size_in_bytes, unit="iB", unit_scale=True, desc="Downloading"
     ) as pb:
         with open(save_path, "wb") as file:
             for data in response.iter_content(block_size):
@@ -61,7 +66,7 @@ def download_file(url: str, save_dir: InputType) -> Path:
     return save_path
 
 
-def unzip_file(file_path: str, save_dir: InputType, is_del_raw: bool = True) -> Path:
+def unzip_file(file_path: str, save_dir: InputType, is_del_raw: bool = True) -> UnzipResult:
     """解压下载得到的tar模型文件，会自动解压到save_dir下以file_path命名的目录下
 
     Args:
@@ -70,15 +75,17 @@ def unzip_file(file_path: str, save_dir: InputType, is_del_raw: bool = True) ->
         is_del_raw (bool, optional): 是否删除原文件. Defaults to True.
 
     Returns:
-        Path: 解压后模型保存路径
+        dict:
+            - model_dir (Path): 解压后模型保存路径
+            - my_files (set): 解压得到的文件名集合
     """
     model_dir = Path(save_dir) / Path(file_path).stem
     mkdir(model_dir)
 
     tar_file_name_list = [".pdiparams", ".pdiparams.info", ".pdmodel", ".json"]
     my_files = set()
     with tarfile.open(file_path, "r") as tarObj:
-        
+
         for member in tarObj.getmembers():
             filename = None
 
@@ -122,4 +129,4 @@ def is_http_url(s: InputType) -> bool:
 
     if regex.match(str(s)):
         return True
-    return False
+    return False