1. Add api subtitle-dl to stream download subtitle.

2. Add id,url, thumbnail to fetchSubtitleUrls, fetchAnySubtitle

fetchYoutubeSubtitle.py

@@ -1,6 +1,7 @@
 import os
 import json
 import math
+import re
 import time
 import traceback
 from typing import Optional
@@ -22,7 +23,52 @@ debug = os.getenv("DEBUG")
 #       }
 #     ]
 #   }
+
+
+def getSubtitleOptions(
+    lang: Optional[str] = None,
+    proxy: Optional[str] = None,
+):
+    ydl_opts = {
+        "noplaylist": True,
+        "writesubtitles": True,
+        "writeautomaticsub": True,
+        # "listsubtitles": True,
+        # "subtitlesformat": subType, # mark due to default youtube no srt and xml format
+        "skip_download": True,
+        "socket_timeout": 10,
+        "extractor_retries": 0,
+        # "debug_printtraffic": True,
+        "extractor_args": {
+            "youtube": {
+                "player_skip": [
+                    "configs",
+                    "initial",
+                ],  # skip "webpage" will cause l2P5PgL1LfI missing some langs,
+                "player_client": ["ios"],
+                "skip": ["hls", "dash"],  # don't skip "translated_subs"
+            }
+        },
+    }
+
+    if lang:
+        ydl_opts.update(
+            {
+                "subtitleslangs": [
+                    lang,
+                    "-live_chat",
+                ]
+            }
+        )  # filter live chat to requested_subtitles
+    if proxy:
+        ydl_opts.update({"proxy": proxy, "socket_timeout": 20})
+
+    # print(ydl_opts)
+    return ydl_opts
+
+
 def getUrlFromSubtitleItem(item, lang="en", subType="vtt"):
+    # print("item: {}, lang: {}, subType: {}".format(item, lang, subType))
     for subtitle in item[lang]:
         if lang != "live_chat" and subType == "xml":
             if debug:
@@ -38,10 +84,21 @@ def getUrlFromSubtitleItem(item, lang="en", subType="vtt"):
     return None
 
 
-def getRequestedSubtitlesUrl(info_dict, lang, subType):
+def getRequestedSubtitlesUrl(info_dict, lang, subType, isLangKey=False):
     item = info_dict.get("requested_subtitles")
-    if item:
-        langs = item.keys()
+    if not item:
+        return None
+
+    langs = item.keys()
+    if lang in langs:
+        item = {lang: [item[lang]]} if type(item[lang]) == dict else item
+        url = getUrlFromSubtitleItem(item, lang, subType)
+        if url:
+            if debug:
+                print("getRequestedSubtitlesUrl lang:{}".format(lang))
+            return url
+
+    if not isLangKey:
         for l in langs:
             if l.startswith(lang):
                 item = {l: [item[l]]} if type(item[l]) == dict else item
@@ -58,6 +115,7 @@ def getSubtitleLangUrl(
     lang="en",
     subType="vtt",
     subTitleKeys=["subtitles", "automatic_captions"],
+    isLangKey=False,
 ):
     for subtitle_item in subTitleKeys:
         langs = info_dict.get(subtitle_item).keys()
@@ -68,15 +126,18 @@ def getSubtitleLangUrl(
                     print("getSubtitleLangUrl lang:{}".format(lang))
                 return url
 
-    for subtitle_item in subTitleKeys:
-        langs = info_dict.get(subtitle_item).keys()
-        for l in langs:
-            if l.startswith(lang):
-                url = getUrlFromSubtitleItem(info_dict.get(subtitle_item), l, subType)
-                if url:
-                    if debug:
-                        print("getSubtitleLangUrl lang:{} url:{}".format(l, url))
-                    return url
+    if not isLangKey:
+        for subtitle_item in subTitleKeys:
+            langs = info_dict.get(subtitle_item).keys()
+            for l in langs:
+                if l.startswith(lang):
+                    url = getUrlFromSubtitleItem(
+                        info_dict.get(subtitle_item), l, subType
+                    )
+                    if url:
+                        if debug:
+                            print("getSubtitleLangUrl lang:{} url:{}".format(l, url))
+                        return url
 
     return None
 
@@ -111,47 +172,21 @@ async def fetchSubtitle(
     subType: Optional[str] = "vtt",
     proxy: Optional[str] = None,
 ) -> dict:
-    return await fetchSubtitlebyType(url, lang, subType, proxy)
+    return await fetchAnySubtitle(url, lang, subType, proxy)
 
 
-async def fetchSubtitlebyType(
+async def fetchAnySubtitle(
     url: str,
     lang: Optional[str] = "en",
     subType: Optional[str] = "vtt",
+    skipEmpty: bool = True,
     proxy: Optional[str] = None,
 ) -> dict:
     # lang-code or lang.* .* is regex
     # reqLang = lang if len(lang.split("-")) > 1 or lang.endswith(".*") else lang + ".*"
 
-    ydl_opts = {
-        "noplaylist": True,
-        "writesubtitles": True,
-        "writeautomaticsub": True,
-        # "listsubtitles": True,
-        # "subtitlesformat": subType, # mark due to default youtube no srt and xml format
-        "subtitleslangs": [
-            lang,
-            "-live_chat",
-        ],  # filter live chat to requested_subtitles
-        "skip_download": True,
-        "socket_timeout": 10,
-        "extractor_retries": 0,
-        # "debug_printtraffic": True,
-        "extractor_args": {
-            "youtube": {
-                "player_skip": [
-                    "configs",
-                    "initial",
-                ],  # skip "webpage" will cause l2P5PgL1LfI missing some langs,
-                "player_client": ["ios"],
-                "skip": ["hls", "dash"],  # don't skip "translated_subs"
-            }
-        },
-    }
+    ydl_opts = getSubtitleOptions(lang, proxy)
 
-    if proxy:
-        ydl_opts.update({"proxy": proxy, "socket_timeout": 20})
-    # print(ydl_opts)
     title = "unknow"
     duration = ""
     try:
@@ -161,52 +196,53 @@ async def fetchSubtitlebyType(
             title = info_dict.get("title", "unknow")
             seconds = info_dict.get("duration")
             duration = str(seconds) if seconds else ""
-            isSrt = False
-            if info_dict.get("extractor") == "youtube" and subType == "srt":
-                subType = "xml"
-                isSrt = True
+            thumbnail = info_dict.get("thumbnail")
+            if ".webp" in thumbnail:
+                thumbnail = "https://i.ytimg.com/vi/{}/hqdefault.jpg".format(
+                    info_dict.get("id")
+                )
+
+            reqType = subType
+            if info_dict.get("extractor") == "youtube" and subType in ["srt", "txt"]:
+                reqType = "xml"
             if debug:
                 print(
                     "subtitles.keys(): {} automatic_captions: {} requested_subtitles: {}".format(
                         info_dict.get("subtitles").keys(),
                         info_dict.get("automatic_captions").keys(),
-                        info_dict.get("requested_subtitles").keys(),
+                        (
+                            info_dict.get("requested_subtitles").keys()
+                            if info_dict.get("requested_subtitles")
+                            else {}
+                        ),
                     )
                 )
 
-            # subtitle_url = getRequestedSubtitlesUrl(info_dict, lang, subType)
-            # if not subtitle_url:
-            #     subtitle_url = getSubtitleLangUrl(info_dict, lang, subType)
-            # if not subtitle_url:
-            #     subtitle_url = getSubtitleOtherUrl(info_dict, lang, subType)
-
             subtitle_funcs = [
                 getRequestedSubtitlesUrl,
                 getSubtitleLangUrl,
                 getSubtitleOtherUrl,
             ]
             for index in range(len(subtitle_funcs)):
-                subtitle_url = subtitle_funcs[index](info_dict, lang, subType)
+                subtitle_url = subtitle_funcs[index](info_dict, lang, reqType)
                 if subtitle_url:
                     # print("subtitle_url: {}".format(subtitle_url))
-                    with ydl.urlopen(subtitle_url) as response:
-                        subtitle = (
-                            xml_caption_to_srt(response.read().decode())
-                            if isSrt
-                            else response.read().decode()
+                    subtitle = fetchSubtitlebydlUrl(ydl, subType, subtitle_url)
+                    print(
+                        "function index:{}, url:{}, title:{}, duration:{} len(subtitle): {}".format(
+                            index, url, title, duration, len(subtitle or "")
                         )
-                        print(
-                            "function index:{}, url:{}, title:{}, duration:{} len(subtitle): {}".format(
-                                index, url, title, duration, len(subtitle or "")
-                            )
-                        )
-                        if subtitle is not None:
-                            return {
-                                "title": title,
-                                "duration": duration,
-                                "subtitle": subtitle,
-                                "chapters": info_dict.get("chapters", None),
-                            }
+                    )
+                    if subtitle is not None:
+                        return {
+                            "id": info_dict.get("id"),
+                            "url": url,
+                            "title": title,
+                            "thumbnail": thumbnail,
+                            "duration": duration,
+                            "subtitle": subtitle,
+                            "chapters": info_dict.get("chapters", None),
+                        }
     except Exception as e:
         print(e)
         traceback.print_exc()
@@ -234,7 +270,7 @@ def is_spaces_only(variable):
     return True
 
 
-def xml_caption_to_srt(xml_captions: str) -> str:
+def xml_caption_to_srt(xml_captions: str, skip_empty: bool = True) -> str:
     """Convert xml caption tracks to "SubRip Subtitle (srt)".
     :param str xml_captions:
         XML formatted caption tracks.
@@ -246,7 +282,7 @@ def xml_caption_to_srt(xml_captions: str) -> str:
         caption = unescape(
             text.replace("\n", " ").replace("  ", " "),
         )
-        if len(caption) == 0 or is_spaces_only(caption):
+        if skip_empty and len(caption) == 0 or is_spaces_only(caption):
             continue
         try:
             duration = float(child.attrib["dur"])
@@ -262,31 +298,39 @@ def xml_caption_to_srt(xml_captions: str) -> str:
             text=caption,
         )
         segments.append(line)
-    # return None if no text in xml
-    return "\n".join(segments).strip() if len(segments) > 0 else None
+
+    if skip_empty:
+        # return None if no text in xml
+        return "\n".join(segments).strip() if len(segments) > 0 else None
+    return "\n".join(segments).strip()
+
+
+def xml_caption_to_txt(xml_captions: str, skip_empty: bool = True) -> str:
+    """Convert xml caption tracks to "SubRip Subtitle (srt)".
+    :param str xml_captions:
+        XML formatted caption tracks.
+    """
+    segments = []
+    root = ElementTree.fromstring(xml_captions)
+    for i, child in enumerate(list(root)):
+        text = child.text or ""
+        caption = unescape(
+            text.replace("\n", " ").replace("  ", " "),
+        )
+        if skip_empty and (len(caption) == 0 or is_spaces_only(caption)):
+            continue
+
+        line = "{text}\n".format(text=caption)
+        segments.append(line)
+
+    if skip_empty:
+        "\n".join(segments).strip() if len(segments) > 0 else None
+
+    return "\n".join(segments).strip()
 
 
 async def fetchSubtitleUrls(url: str, proxy: Optional[str] = None) -> json:
-    ydl_opts = {
-        "noplaylist": True,
-        "writesubtitles": True,
-        "writeautomaticsub": True,
-        # "allsubtitles": False,
-        # "listsubtitles": True,
-        "skip_download": True,
-        "socket_timeout": 10,
-        "extractor_retries": 0,
-        # "debug_printtraffic": True,
-        "extractor_args": {
-            "youtube": {
-                "player_skip": ["configs", "initial"],
-                "player_client": ["ios"],
-                "skip": ["hls", "dash"],  # , "translated_subs"
-            }
-        },
-    }
-    if proxy:
-        ydl_opts.update({"proxy": proxy, "socket_timeout": 20})
+    ydl_opts = getSubtitleOptions(proxy)
 
     title = "unknow"
     duration = ""
@@ -296,9 +340,16 @@ async def fetchSubtitleUrls(url: str, proxy: Optional[str] = None) -> json:
             title = info_dict.get("title", "unknow")
             seconds = info_dict.get("duration")
             duration = str(seconds) if seconds else ""
-
+            thumbnail = info_dict.get("thumbnail")
+            if ".webp" in thumbnail:
+                thumbnail = "https://i.ytimg.com/vi/{}/hqdefault.jpg".format(
+                    info_dict.get("id")
+                )
             return {
+                "id": info_dict.get("id"),
+                "url": url,
                 "title": title,
+                "thumbnail": thumbnail,
                 "duration": duration,
                 "subtitles": info_dict.get("subtitles"),
                 "automatic_captions": info_dict.get("automatic_captions"),
@@ -306,3 +357,81 @@ async def fetchSubtitleUrls(url: str, proxy: Optional[str] = None) -> json:
 
     except Exception as e:
         return {"error": str(e)}
+
+
+def fetchSubtitlebydlUrl(ydl, subType, dlUrl, skipEmpty=True):
+    dlUrl = dlUrl if subType not in ["srt", "txt"] else re.sub(r"&fmt=[\w]+", "", dlUrl)
+
+    try:
+        with ydl.urlopen(dlUrl) as resp:
+            if subType == "srt":
+                return xml_caption_to_srt(resp.read().decode(), skipEmpty)
+            elif subType == "txt":
+                return xml_caption_to_txt(resp.read().decode(), skipEmpty)
+            else:
+                return resp.read().decode()
+    except Exception as e:
+        print(e)
+        return None
+
+
+def getSubtitleUrlByLang(info_dict, lang, subType, isLangKey):
+    subtitle_funcs = [
+        getRequestedSubtitlesUrl,
+        getSubtitleLangUrl,
+    ]
+    for index in range(len(subtitle_funcs)):
+        subtitle_url = subtitle_funcs[index](
+            info_dict, lang, subType, isLangKey=isLangKey
+        )
+        print("getSubtitleUrlByLang subtitle_url: {}".format(subtitle_url))
+        if subtitle_url:
+            return subtitle_url
+
+
+async def fetchSubtitleByInfo(
+    url: str, subType: str, dlInfo, proxy: Optional[str] = None
+):
+    try:
+        reqType = "xml" if subType in ["srt", "txt"] else subType
+        ydl_opts = getSubtitleOptions(dlInfo.get("lang", None), proxy)
+        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
+            subtitle = None
+            if "dlUrl" in dlInfo:
+                subtitle = fetchSubtitlebydlUrl(
+                    ydl, subType, dlInfo.get("dlUrl"), False
+                )
+                if subtitle is not None:
+                    return subtitle
+
+            info_dict = ydl.extract_info(url, download=False)
+            if debug:
+                print(
+                    "subtitles.keys(): {} automatic_captions: {} requested_subtitles: {}".format(
+                        info_dict.get("subtitles").keys(),
+                        info_dict.get("automatic_captions").keys(),
+                        (
+                            info_dict.get("requested_subtitles").keys()
+                            if info_dict.get("requested_subtitles")
+                            else {}
+                        ),
+                    )
+                )
+
+            subtitleUrl = None
+            if "langKey" in dlInfo:
+                subtitleUrl = getSubtitleUrlByLang(
+                    info_dict, dlInfo.get("langKey"), reqType, True
+                )
+            if subtitleUrl is None:
+                subtitleUrl = getSubtitleUrlByLang(
+                    info_dict, dlInfo.get("lang"), reqType, False
+                )
+
+            print("subtitleUrl: {}".format(subtitleUrl))
+            subtitle = fetchSubtitlebydlUrl(ydl, subType, subtitleUrl, False)
+            return subtitle
+    except Exception as e:
+        print(e)
+        traceback.print_exc()
+        return {"error": str(e)}

main.py

@@ -1,14 +1,17 @@
+import json
 import os
-from fastapi import FastAPI, Header, Request
-from fastapi.responses import JSONResponse
-from typing import Annotated
-from fetchYoutubeSubtitle import fetchSubtitle, fetchSubtitleUrls
+from fastapi import FastAPI, Header, Request, Response, HTTPException
+from fastapi.responses import JSONResponse, StreamingResponse
+import traceback
+from typing import Annotated, Optional
+from fetchYoutubeSubtitle import fetchSubtitle, fetchSubtitleUrls, fetchSubtitleByInfo
 
 token = os.getenv("HF_X_TOKEN")
 app = FastAPI()
 
+
 @app.get("/")
-def read_root(request:Request):
+def read_root(request: Request):
     print(request.headers)
     print(request.client.host)
     print(request.client.port)
@@ -21,7 +24,13 @@ def read_json():
 
 
 @app.get("/subtitle/")
-async def get_subtitle(url: str, subtype: str="srt", lang:str='en', proxy: str=None, x_token: Annotated[str | None, Header()] = None):
+async def get_subtitle(
+    url: str,
+    subtype: str = "srt",
+    lang: str = "en",
+    proxy: str = None,
+    x_token: Annotated[str | None, Header()] = None,
+):
     if token != x_token:
         return JSONResponse({"error": "Invalid token"})
     subtitle = await fetchSubtitle(url, lang=lang, subType=subtype, proxy=proxy)
@@ -29,8 +38,68 @@ async def get_subtitle(url: str, subtype: str="srt", lang:str='en', proxy: str=N
 
 
 @app.get("/subtitle-urls/")
-async def get_subtitleUrls(url: str, proxy: str=None, x_token: Annotated[str | None, Header()] = None):
+async def get_subtitleUrls(
+    url: str, proxy: str = None, x_token: Annotated[str | None, Header()] = None
+):
     if token != x_token:
         return JSONResponse({"error": "Invalid token"})
     subtitles = await fetchSubtitleUrls(url, proxy=proxy)
-    return JSONResponse(content=subtitles)
+    return JSONResponse(content=subtitles)
+
+
+def download_file(content, chunk_size):
+    num_chunks = (len(content) // chunk_size) + 1
+    for i in range(num_chunks):
+        start = i * chunk_size
+        end = (i + 1) * chunk_size
+        yield content[start:end]
+
+
+@app.get("/subtitle-dl/")
+async def download(
+    url: str,
+    fileName: str,
+    fileType: str,
+    info: str,  # download info
+    proxy: str = None,
+    x_token: Annotated[str | None, Header()] = None,
+    request: Request = None,
+):
+    if token != x_token:
+        raise HTTPException(status_code=401, detail="Invalid token")
+
+    try:
+        dlInfo = json.loads(info)
+        # print(
+        #     "url: {}, fileName: {}, fileType: {}, dlInfo: {}".format(
+        #         url, fileName, fileType, dlInfo
+        #     )
+        # )
+        subtitle = await fetchSubtitleByInfo(url, fileType, dlInfo, proxy=proxy)
+
+        excluded_headers = [
+            "content-encoding",
+            "content-length",
+            "transfer-encoding",
+            "connection",
+        ]
+        headers = [
+            (name, value)
+            for (name, value) in request.headers.items()
+            if name.lower() not in excluded_headers
+        ]
+        headers.append(("Content-Type", "text/plain"))
+        headers.append(
+            (
+                "Content-Disposition",
+                f'attachment; filename="{fileName.encode("utf-8").decode("latin-1")}.{fileType}"',
+            )
+        )
+
+        return StreamingResponse(
+            download_file(subtitle, 8192), headers=dict(headers), status_code=200
+        )
+    except Exception as e:
+        print(e)
+        traceback.print_exc()
+        raise HTTPException(status_code=500, detail="Internal Server Error")