tiktok.py

import os
import re
import shutil
import requests
import gradio as gr
from utils import timestamp
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)
    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:
            for chunk in response.iter_content(chunk_size=8192):
                file.write(chunk)

    print(f"[{timestamp()}] File was downloaded to {local_file}")
    return local_file


def extract_fst_url(text):
    url_pattern = r'(https?://[^\s"]+)'
    match = re.search(url_pattern, text)
    if match:
        return match.group(1)
    else:
        return None


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:
        response = requests.get(API_TIKTOK, params={"url": video_url}, timeout=TIMEOUT)
        response_json = response.json()
        retcode = response_json["code"]
        if retcode == 200:
            response_data = response_json["data"]
            video_id = response_data["play_url"].split("video_id=")[1].split("&")[0]
            video = download_file(response_data["video_url"], video_id)
            parse_time = response_data["parse_time"]

            additional_data = response_data["additional_data"][0]
            desc = additional_data["desc"]
            avatar = additional_data["url"].split("?from=")[0]
            author = additional_data["nickname"]
            sign = additional_data["signature"]

        else:
            desc = f"Interface call failed, error code: HTTP {retcode}"

    except Exception as e:
        desc = f"Video parsing failed: {e}"

    return video, desc, parse_time, avatar, author, sign


def tiktok_parser():
    return gr.Interface(
        fn=infer,
        inputs=[
            gr.Textbox(
                label="Please enter TikTok video sharing short link",
                placeholder="https://v.douyin.com/*",
            ),
        ],
        outputs=[
            gr.Video(
                label="Video download",
                show_download_button=True,
                show_share_button=False,
            ),
            gr.Textbox(label="Video description", show_copy_button=True),
            gr.Textbox(label="Parsing time", show_copy_button=True),
            gr.Image(label="Author avatar", show_share_button=False),
            gr.Textbox(label="Author nickname", show_copy_button=True),
            gr.TextArea(label="Author signature", show_copy_button=True),
        ],
        title="Parse TikTok video without watermark",
        flagging_mode="never",
        examples=[
            "https://v.douyin.com/8FVe5DzarE0",
            "8.20 Njc:/ [email protected] 11/03 黑塔女士举世无双！# 大黑塔 # 黑塔 # 崩坏星穹铁道 # 再创世的凯歌 # 天才俱乐部  https://v.douyin.com/8FVe5DzarE0/ 复制此链接，打开Dou音搜索，直接观看视频！",
        ],
        cache_examples=False,
    )