Add support for translating vtt files

Author: CoolCat467

.pre-commit-config.yaml

@@ -23,7 +23,7 @@ repos:
     hooks:
       - id: black
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.9.6
+    rev: v0.9.7
     hooks:
       - id: ruff
         types: [file]
@@ -40,7 +40,7 @@ repos:
         additional_dependencies:
           - tomli
   - repo: https://github.com/crate-ci/typos
-    rev: typos-dict-v0.12.5
+    rev: v1.29.9
     hooks:
       - id: typos
   - repo: https://github.com/woodruffw/zizmor-pre-commit

src/subtitle_translate/main.py

@@ -27,12 +27,17 @@
 
 
 import argparse
+import sys
+from typing import TYPE_CHECKING
 
 import httpx
 import trio
 
 from subtitle_translate import extricate, subtitle_parser, translate
 
+if TYPE_CHECKING:
+    from collections.abc import Iterable
+
 
 async def translate_texts(
     texts: dict[int, tuple[str, ...]],
@@ -61,16 +66,12 @@ async def translate_texts(
 
 
 async def translate_subtitles(
-    source_file: str,
-    dest_file: str,
+    generator: Iterable[tuple[int, subtitle_parser.Subtitle]],
     source_language: str = "auto",
     dest_language: str = "en",
-) -> None:
+) -> tuple[dict[int, subtitle_parser.Subtitle], dict[int, tuple[str, ...]]]:
     """Translate subtitles file asynchronously."""
-    print(f"Loading subtitles file {source_file!r}...")
-    subs, texts = subtitle_parser.convert_text(
-        subtitle_parser.parse_file(source_file),
-    )
+    subs, texts = subtitle_parser.convert_text(generator)
 
     print(f"Parsed {len(subs)} subtitles")
 
@@ -80,11 +81,76 @@ async def translate_subtitles(
     sentence_count = sum(map(len, texts.values()))
     print(f"Translated {sentence_count} sentences.")
 
+    return subs, new_texts
+
+
+async def translate_subtitles_srt(
+    source_file: str,
+    dest_file: str | None = None,
+    source_language: str = "auto",
+    dest_language: str = "en",
+) -> None:
+    """Translate subtitles file asynchronously."""
+    # Set destination if not provided
+    if dest_file is None:
+        name, ext = source_file.rsplit(".", 1)
+        dest_file = f"{name}.{dest_language}.{ext}"
+
+    print(f"Loading subtitles file {source_file!r}...")
+    subs, new_texts = await translate_subtitles(
+        subtitle_parser.parse_file_srt(source_file),
+    )
+
     print("Updating subtitle texts...")
     subs = subtitle_parser.modify_subtitles(subs, new_texts)
 
     print("Saving...")
-    subtitle_parser.write_subtitles_file(dest_file, subs)
+    subtitle_parser.write_subtitles_srt_file(dest_file, subs)
+    print("Save complete.")
+    print(f"Saved to {dest_file!r}")
+
+
+async def translate_subtitles_vtt(
+    source_file: str,
+    dest_file: str | None = None,
+    source_language: str = "auto",
+    dest_language: str = "en",
+) -> None:
+    """Translate subtitles file asynchronously."""
+    print(f"Loading subtitles file {source_file!r}...")
+    gen = subtitle_parser.parse_file_vtt(source_file)
+    header = next(gen)
+
+    new_header = []
+    for line in header:
+        assert isinstance(line, str)
+        if source_language == "auto" and line.startswith("Language: "):
+            key, value = line.split(" ", 1)
+            source_language = value
+            line = f"{key} {dest_language}"
+        new_header.append(line)
+
+    # Set destination if not provided
+    if dest_file is None:
+        name, ext = source_file.rsplit(".", 1)
+        source_lang_ext = f".{source_language}"
+        if source_lang_ext in name:
+            name = name.removesuffix(source_lang_ext)
+        dest_file = f"{name}.{dest_language}.{ext}"
+
+    subs, new_texts = await translate_subtitles(
+        enumerate(x for x in gen if isinstance(x, subtitle_parser.Subtitle)),
+    )
+
+    print("Updating subtitle texts...")
+    subs = subtitle_parser.modify_subtitles_plain(subs, new_texts)
+
+    print("Saving...")
+    subtitle_parser.write_subtitles_vtt_file(
+        dest_file,
+        new_header,
+        subs.values(),
+    )
     print("Save complete.")
     print(f"Saved to {dest_file!r}")
 
@@ -119,6 +185,15 @@ async def run_async() -> None:
             "Must be a ISO 639-1:2002 language code or 'auto' to guess."
         ),
     )
+    parser.add_argument(
+        "--source-type",
+        type=str,
+        default="auto",
+        help=(
+            "Subtitle source type (default: 'auto').\n"
+            "Must be either 'srt' or 'vtt', or 'auto' to guess from filename."
+        ),
+    )
     parser.add_argument(
         "--dest-lang",
         type=str,
@@ -136,17 +211,27 @@ async def run_async() -> None:
 
     args = parser.parse_args()
 
-    # Set destination if not provided
-    if args.dest_file is None:
+    if args.source_type == "auto":
         name, ext = args.source_file.rsplit(".", 1)
-        args.dest_file = f"{name}.{args.dest_lang}.{ext}"
-
-    await translate_subtitles(
-        args.source_file,
-        args.dest_file,
-        args.source_lang,
-        args.dest_lang,
-    )
+        args.source_type = ext
+
+    if args.source_type == "srt":
+        await translate_subtitles_srt(
+            args.source_file,
+            args.dest_file,
+            args.source_lang,
+            args.dest_lang,
+        )
+    elif args.source_type == "vtt":
+        await translate_subtitles_vtt(
+            args.source_file,
+            args.dest_file,
+            args.source_lang,
+            args.dest_lang,
+        )
+    else:
+        print(f"Unhandled source type {args.source_type!r}.")
+        sys.exit(1)
 
 
 def cli_run() -> None: