sync ms

bili.py

@@ -1,25 +1,30 @@
-import os
 import re
-import shutil
 import requests
 import gradio as gr
-from utils import timestamp
+from tqdm import tqdm
+from utils import timestamp, clean_dir
 from config import TMP_DIR, HEADER, TIMEOUT, API_BILI
 
 
 def download_file(url, video_id, cache_dir=TMP_DIR):
-    if os.path.exists(cache_dir):
-        shutil.rmtree(cache_dir)
-
-    os.makedirs(cache_dir)
+    clean_dir(cache_dir)
     local_file = f"{cache_dir}/{video_id}.mp4"
     response = requests.get(url, stream=True)
     if response.status_code == 200:
-        with open(local_file, "wb") as file:
+        total_size = int(response.headers.get("Content-Length", 0)) + 1
+        time_stamp = timestamp()
+        progress_bar = tqdm(
+            total=total_size,
+            unit="B",
+            unit_scale=True,
+            desc=f"[{time_stamp}] {local_file}",
+        )
+        with open(local_file, "wb") as f:
             for chunk in response.iter_content(chunk_size=8192):
-                file.write(chunk)
+                if chunk:
+                    f.write(chunk)
+                    progress_bar.update(len(chunk))
 
-    print(f"[{timestamp()}] File was downloaded to {local_file}")
     return local_file
 
 
@@ -46,14 +51,15 @@ def get_real_url(short_url):
     ).url.split("/?")[0]
 
 
+# outer func
 def infer(video_url: str, p: int):
     title = cover = desc = dur = video = author = avatar = None
     if not video_url:
         title = "Empty video link!"
         return title, cover, video, desc, dur, avatar, author
 
-    video_url = extract_fst_url(video_url)
     try:
+        video_url = extract_fst_url(video_url)
         if "b23.tv" in video_url:
             video_url = get_real_url(video_url)
 
@@ -75,7 +81,7 @@ def infer(video_url: str, p: int):
             avatar = author_data["user_img"]
 
         else:
-            title = f"Failed to call API, error code: {retcode}"
+            raise ConnectionError(f"Failed to call API, error code: {retcode}")
 
     except Exception as e:
         title = f"Failed to parse video: {e}"

bvid2acid.py

@@ -2,6 +2,7 @@ import requests
 import gradio as gr
 
 
+# outer func
 def infer(bvid: str):
     try:
         response = requests.get(

tiktok.py

@@ -1,25 +1,30 @@
-import os
 import re
-import shutil
 import requests
 import gradio as gr
-from utils import timestamp
+from tqdm import tqdm
+from utils import timestamp, clean_dir
 from config import API_TIKTOK, TIMEOUT, TMP_DIR
 
 
-def download_file(url, video_id, cache_dir=TMP_DIR):
-    if os.path.exists(cache_dir):
-        shutil.rmtree(cache_dir)
-
-    os.makedirs(cache_dir)
+def download_file(url, video_id, cache_dir=f"{TMP_DIR}/tiktok"):
+    clean_dir(cache_dir)
     local_file = f"{cache_dir}/{video_id}.mp4"
     response = requests.get(url, stream=True)
     if response.status_code == 200:
-        with open(local_file, "wb") as file:
+        total_size = int(response.headers.get("Content-Length", 0)) + 1
+        time_stamp = timestamp()
+        progress_bar = tqdm(
+            total=total_size,
+            unit="B",
+            unit_scale=True,
+            desc=f"[{time_stamp}] {local_file}",
+        )
+        with open(local_file, "wb") as f:
             for chunk in response.iter_content(chunk_size=8192):
-                file.write(chunk)
+                if chunk:
+                    f.write(chunk)
+                    progress_bar.update(len(chunk))
 
-    print(f"[{timestamp()}] File was downloaded to {local_file}")
     return local_file
 
 
@@ -32,18 +37,18 @@ def extract_fst_url(text):
         return None
 
 
+# outer func
 def infer(video_url):
     video = parse_time = desc = avatar = author = sign = None
     if not video_url:
         desc = "The video sharing link is empty!"
         return video, desc, parse_time, avatar, author, sign
 
-    video_url = extract_fst_url(video_url)
-    if not video_url:
-        desc = "Please enter a valid video sharing link!"
-        return video, desc, parse_time, avatar, author, sign
-
     try:
+        video_url = extract_fst_url(video_url)
+        if not video_url:
+            raise ValueError("Please enter a valid video sharing link!")
+
         response = requests.get(API_TIKTOK, params={"url": video_url}, timeout=TIMEOUT)
         response_json = response.json()
         retcode = response_json["code"]
@@ -60,7 +65,7 @@ def infer(video_url):
             sign = additional_data["signature"]
 
         else:
-            desc = f"Interface call failed, error code: HTTP {retcode}"
+            raise ConnectionError(f"Interface call failed, error code: HTTP {retcode}")
 
     except Exception as e:
         desc = f"Video parsing failed: {e}"

utils.py

@@ -1,3 +1,5 @@
+import os
+import shutil
 from datetime import datetime
 from zoneinfo import ZoneInfo
 from tzlocal import get_localzone
@@ -10,3 +12,18 @@ def timestamp(naive_time: datetime = None, target_tz=ZoneInfo("Asia/Shanghai")):
     local_tz = get_localzone()
     aware_local = naive_time.replace(tzinfo=local_tz)
     return aware_local.astimezone(target_tz).strftime("%Y-%m-%d %H:%M:%S")
+
+
+def mk_dir(dirpath: str):
+    if not os.path.exists(dirpath):
+        os.makedirs(dirpath)
+
+
+def rm_dir(dirpath: str):
+    if os.path.exists(dirpath):
+        shutil.rmtree(dirpath)
+
+
+def clean_dir(dirpath: str):
+    rm_dir(dirpath)
+    os.makedirs(dirpath)