|
| 1 | +"""Run logic for `assembly caption`: transcribe → SRT export → ffmpeg burn-in. |
| 2 | +
|
| 3 | +The command module (aai_cli/commands/caption.py) only parses argv — it builds a |
| 4 | +``CaptionOptions`` and hands it to ``run_caption`` via ``context.run_command`` |
| 5 | +(the options/run split, see AGENTS.md), so tests drive the whole pipeline by |
| 6 | +constructing options directly. |
| 7 | +
|
| 8 | +The pipeline: the video is transcribed (or an existing transcript is reused via |
| 9 | +``--transcript-id``), the transcript's SRT captions are fetched from the export |
| 10 | +endpoint, and ffmpeg's ``subtitles`` filter burns them into the picture (open |
| 11 | +captions, always visible) while the audio stream is copied untouched. A |
| 12 | +YouTube/media-page URL is downloaded first — always the full video, since the |
| 13 | +captions are burned into it. |
| 14 | +""" |
| 15 | + |
| 16 | +from __future__ import annotations |
| 17 | + |
| 18 | +import shutil |
| 19 | +import subprocess |
| 20 | +import tempfile |
| 21 | +from dataclasses import dataclass |
| 22 | +from pathlib import Path |
| 23 | + |
| 24 | +import assemblyai as aai |
| 25 | +from rich.markup import escape |
| 26 | + |
| 27 | +from aai_cli import client, output, youtube |
| 28 | +from aai_cli.context import AppState |
| 29 | +from aai_cli.errors import CLIError, UsageError |
| 30 | + |
| 31 | + |
| 32 | +@dataclass(frozen=True) |
| 33 | +class CaptionOptions: |
| 34 | + """Every `assembly caption` flag as plain data (``--json`` excluded: |
| 35 | + run_command resolves it into the ``json_mode`` argument).""" |
| 36 | + |
| 37 | + # The raw source as typed: a local path, or a downloadable media-page URL |
| 38 | + # (a pathlib.Path would collapse the "//" in "https://"). |
| 39 | + media: str |
| 40 | + transcript_id: str | None |
| 41 | + chars_per_caption: int | None |
| 42 | + font_size: int | None |
| 43 | + out: Path | None |
| 44 | + |
| 45 | + |
| 46 | +def default_out_path(media: Path) -> Path: |
| 47 | + """The default output file: ``<stem>.captioned<ext>`` next to the input.""" |
| 48 | + return media.parent / f"{media.stem}.captioned{media.suffix}" |
| 49 | + |
| 50 | + |
| 51 | +# ffmpeg's filtergraph syntax gives these characters meaning (option/filter/chain |
| 52 | +# separators, stream labels, quoting), so a path embedded in `-vf subtitles=…` |
| 53 | +# must escape them or a TMPDIR containing one would corrupt the filter spec. |
| 54 | +_FILTER_ESCAPES = str.maketrans({ch: f"\\{ch}" for ch in "\\':,;[]"}) |
| 55 | + |
| 56 | + |
| 57 | +def subtitles_filter(srt: Path, font_size: int | None) -> str: |
| 58 | + """The ``-vf`` filtergraph burning ``srt`` into the video.""" |
| 59 | + spec = f"subtitles={str(srt).translate(_FILTER_ESCAPES)}" |
| 60 | + if font_size is not None: |
| 61 | + spec += f":force_style=FontSize={font_size}" |
| 62 | + return spec |
| 63 | + |
| 64 | + |
| 65 | +def _validate_media(media: Path) -> None: |
| 66 | + """Reject a missing local source before credential resolution, so a typo'd |
| 67 | + path reads as "file not found", never as a login prompt or an ffmpeg error.""" |
| 68 | + if not media.exists(): |
| 69 | + raise CLIError( |
| 70 | + f"File not found: {media}", |
| 71 | + error_type="file_not_found", |
| 72 | + exit_code=2, |
| 73 | + suggestion="Check the path. assembly caption needs a local video file.", |
| 74 | + ) |
| 75 | + if not media.is_file(): |
| 76 | + raise CLIError( |
| 77 | + f"Not a file: {media}", |
| 78 | + error_type="not_a_file", |
| 79 | + exit_code=2, |
| 80 | + suggestion="Pass a video file, not a directory.", |
| 81 | + ) |
| 82 | + |
| 83 | + |
| 84 | +def _validate_out(out: Path, media: Path) -> None: |
| 85 | + """The captioned file must never overwrite its own input: ffmpeg would read |
| 86 | + and write the same file concurrently, corrupting it.""" |
| 87 | + if out.resolve() == media.resolve(): |
| 88 | + raise UsageError( |
| 89 | + "--out would overwrite the input file.", |
| 90 | + suggestion="Pick a different output path.", |
| 91 | + ) |
| 92 | + |
| 93 | + |
| 94 | +def _require_ffmpeg() -> str: |
| 95 | + """The ffmpeg executable; checked before any (billed) transcription work.""" |
| 96 | + path = shutil.which("ffmpeg") |
| 97 | + if path is None: |
| 98 | + raise CLIError( |
| 99 | + "ffmpeg is required to burn captions into video, but it isn't on PATH.", |
| 100 | + error_type="missing_dependency", |
| 101 | + suggestion="Install it (brew install ffmpeg / apt install ffmpeg) and re-run.", |
| 102 | + ) |
| 103 | + return path |
| 104 | + |
| 105 | + |
| 106 | +def _run_ffmpeg(args: list[str]) -> subprocess.CompletedProcess[str]: |
| 107 | + """Boundary seam for tests: one ffmpeg invocation, output captured.""" |
| 108 | + return subprocess.run(args, capture_output=True, text=True, check=False) |
| 109 | + |
| 110 | + |
| 111 | +def _burn(ffmpeg: str, media: Path, srt: Path, out: Path, font_size: int | None) -> None: |
| 112 | + """Burn the ``srt`` captions into ``media``'s video stream, writing ``out``. |
| 113 | +
|
| 114 | + The video is necessarily re-encoded (the captions become pixels); ``-c:a |
| 115 | + copy`` carries the audio over untouched. The explicit ``-map 0:v`` makes |
| 116 | + audio-only input an ffmpeg error ("matches no streams") instead of a silent |
| 117 | + uncaptioned copy; ``-map 0:a?`` keeps a silent video legal. ``-y`` makes a |
| 118 | + re-run overwrite its own earlier output instead of stalling on ffmpeg's |
| 119 | + prompt. |
| 120 | + """ |
| 121 | + result = _run_ffmpeg( |
| 122 | + [ |
| 123 | + ffmpeg, |
| 124 | + "-hide_banner", |
| 125 | + "-loglevel", |
| 126 | + "error", |
| 127 | + "-y", |
| 128 | + "-i", |
| 129 | + str(media), |
| 130 | + "-vf", |
| 131 | + subtitles_filter(srt, font_size), |
| 132 | + "-map", |
| 133 | + "0:v", |
| 134 | + "-map", |
| 135 | + "0:a?", |
| 136 | + "-c:a", |
| 137 | + "copy", |
| 138 | + str(out), |
| 139 | + ] |
| 140 | + ) |
| 141 | + if result.returncode != 0: |
| 142 | + detail = result.stderr.strip().splitlines() |
| 143 | + reason = detail[-1] if detail else f"ffmpeg exited with code {result.returncode}" |
| 144 | + raise CLIError( |
| 145 | + f"Could not write {out.name}: {reason}", |
| 146 | + error_type="caption_failed", |
| 147 | + suggestion="Check that the input is a readable video file — captions " |
| 148 | + "can't be burned into audio-only media.", |
| 149 | + ) |
| 150 | + |
| 151 | + |
| 152 | +def _resolve_transcript( |
| 153 | + opts: CaptionOptions, media: Path, state: AppState, *, json_mode: bool |
| 154 | +) -> object: |
| 155 | + """The transcript whose captions are burned in: fetched by id, or made fresh |
| 156 | + from the (already local) media file.""" |
| 157 | + if opts.transcript_id is not None: |
| 158 | + return client.get_transcript(state.resolve_api_key(), opts.transcript_id) |
| 159 | + api_key = state.resolve_api_key() |
| 160 | + with output.status("Transcribing for captions…", json_mode=json_mode, quiet=state.quiet): |
| 161 | + return client.transcribe(api_key, str(media), config=aai.TranscriptionConfig()) |
| 162 | + |
| 163 | + |
| 164 | +def _fetch_srt(transcript: object, opts: CaptionOptions, *, json_mode: bool, quiet: bool) -> str: |
| 165 | + """The transcript's SRT captions from the export endpoint; empty is an error.""" |
| 166 | + with output.status("Fetching captions…", json_mode=json_mode, quiet=quiet): |
| 167 | + srt = client.select_transcript_field( |
| 168 | + transcript, "srt", chars_per_caption=opts.chars_per_caption |
| 169 | + ) |
| 170 | + if not srt.strip(): |
| 171 | + transcript_id = str(getattr(transcript, "id", "")) |
| 172 | + raise CLIError( |
| 173 | + f"Transcript {transcript_id} has no captions to burn in.", |
| 174 | + error_type="no_captions", |
| 175 | + exit_code=2, |
| 176 | + suggestion="The media may contain no speech; check it with " |
| 177 | + "'assembly transcribe <file>'.", |
| 178 | + ) |
| 179 | + return srt |
| 180 | + |
| 181 | + |
| 182 | +def run_caption(opts: CaptionOptions, state: AppState, *, json_mode: bool) -> None: |
| 183 | + """Execute one `assembly caption` invocation from already-parsed flags.""" |
| 184 | + ffmpeg = _require_ffmpeg() |
| 185 | + if youtube.is_downloadable_url(opts.media): |
| 186 | + # A media-page URL (YouTube, …) is downloaded once — always the full |
| 187 | + # video, since the captions are burned into it. The download dir is |
| 188 | + # temporary, so the default output lands in the current directory. |
| 189 | + with tempfile.TemporaryDirectory(prefix="aai-caption-src-") as td: |
| 190 | + with output.status("Downloading video…", json_mode=json_mode, quiet=state.quiet): |
| 191 | + local = youtube.download_media(opts.media, Path(td), video=True) |
| 192 | + out = opts.out if opts.out is not None else Path.cwd() / default_out_path(local).name |
| 193 | + _validate_out(out, local) |
| 194 | + _caption_and_emit(opts, local, out, ffmpeg, state, json_mode=json_mode) |
| 195 | + return |
| 196 | + if opts.media.startswith(("http://", "https://")): |
| 197 | + raise UsageError( |
| 198 | + "assembly caption can't fetch this URL; it captions a local file or a " |
| 199 | + "media-page URL yt-dlp can download (YouTube, …).", |
| 200 | + suggestion="Download the video first, then caption the local copy.", |
| 201 | + ) |
| 202 | + media = Path(opts.media) |
| 203 | + _validate_media(media) |
| 204 | + out = opts.out if opts.out is not None else default_out_path(media) |
| 205 | + _validate_out(out, media) |
| 206 | + _caption_and_emit(opts, media, out, ffmpeg, state, json_mode=json_mode) |
| 207 | + |
| 208 | + |
| 209 | +def _caption_and_emit( |
| 210 | + opts: CaptionOptions, |
| 211 | + media: Path, |
| 212 | + out: Path, |
| 213 | + ffmpeg: str, |
| 214 | + state: AppState, |
| 215 | + *, |
| 216 | + json_mode: bool, |
| 217 | +) -> None: |
| 218 | + """Caption an already-local video file into ``out`` and report the result.""" |
| 219 | + transcript = _resolve_transcript(opts, media, state, json_mode=json_mode) |
| 220 | + transcript_id = str(getattr(transcript, "id", "")) |
| 221 | + srt = _fetch_srt(transcript, opts, json_mode=json_mode, quiet=state.quiet) |
| 222 | + captions = srt.count("-->") # one arrow per SRT cue timing line |
| 223 | + with tempfile.TemporaryDirectory(prefix="aai-caption-") as tmp: |
| 224 | + srt_path = Path(tmp) / "captions.srt" |
| 225 | + srt_path.write_text(srt, encoding="utf-8") |
| 226 | + with output.status("Burning captions…", json_mode=json_mode, quiet=state.quiet): |
| 227 | + _burn(ffmpeg, media, srt_path, out, opts.font_size) |
| 228 | + payload: dict[str, object] = { |
| 229 | + "source": opts.media, |
| 230 | + "out": str(out), |
| 231 | + "transcript_id": transcript_id, |
| 232 | + "captions": captions, |
| 233 | + } |
| 234 | + output.emit( |
| 235 | + payload, |
| 236 | + lambda _: output.success(f"{escape(str(out))} {captions} caption(s) burned in"), |
| 237 | + json_mode=json_mode, |
| 238 | + ) |
0 commit comments