From cc4c95e0d283320b4159794b6f78deb71496f156 Mon Sep 17 00:00:00 2001 From: Alex Kroman Date: Thu, 11 Jun 2026 09:23:24 -0700 Subject: [PATCH] Add `aai transcribe --download-sections` (yt-dlp passthrough) Wire yt-dlp's `--download-sections` flag through to `aai transcribe` so a YouTube/podcast URL can fetch only part of the source (e.g. `*0:00-5:00` for the first five minutes) before transcribing. - youtube.download_audio gains a `download_sections` arg that sets yt-dlp's `download_ranges` + `force_keyframes_at_cuts` (exact cuts). A new parse_download_sections() mirrors yt-dlp's grammar verbatim: `*start-end` timestamp ranges (comma-separated, inf/open-ended/negative bounds), chapter regexes, and `*from-url`. Malformed specs raise a clean UsageError (exit 2). - transcribe command exposes the repeatable `--download-sections` option, threaded through run_transcription and reflected in `--show-code`. - code_gen renders the sections into the generated yt-dlp block (download_range_func, force_keyframes_at_cuts, conditional `import re`). - Moved the transcribe `validate_*` helpers into transcribe_exec.py to keep commands/transcribe.py under the 500-line gate. Co-Authored-By: Claude Opus 4.8 (1M context) --- aai_cli/code_gen/__init__.py | 5 +- aai_cli/code_gen/transcribe.py | 82 +++++++++++-- aai_cli/commands/transcribe.py | 59 ++++------ aai_cli/transcribe_exec.py | 32 ++++- aai_cli/youtube.py | 90 +++++++++++++- scripts/generated_code_compile_gate.py | 10 ++ .../test_cli_output_snapshots.ambr | 8 ++ tests/test_code_gen.py | 36 ++++++ tests/test_transcribe.py | 52 +++++++- tests/test_youtube.py | 111 +++++++++++++++++- 10 files changed, 434 insertions(+), 51 deletions(-) diff --git a/aai_cli/code_gen/__init__.py b/aai_cli/code_gen/__init__.py index 7758414d..cd5f5f96 100644 --- a/aai_cli/code_gen/__init__.py +++ b/aai_cli/code_gen/__init__.py @@ -33,9 +33,12 @@ def transcribe( source: str, *, llm_gateway: dict[str, object] | None = None, + download_sections: list[str] | None = None, ) -> str: """Generate runnable Python that reproduces this transcribe invocation.""" - return _transcribe.render(merged, source, llm_gateway=llm_gateway) + return _transcribe.render( + merged, source, llm_gateway=llm_gateway, download_sections=download_sections + ) def stream( diff --git a/aai_cli/code_gen/transcribe.py b/aai_cli/code_gen/transcribe.py index fe90a953..403a53e8 100644 --- a/aai_cli/code_gen/transcribe.py +++ b/aai_cli/code_gen/transcribe.py @@ -26,6 +26,7 @@ def render( *, llm_gateway: dict[str, object] | None = None, output: str | None = None, + download_sections: list[str] | None = None, ) -> str: """Generate a runnable transcribe script reproducing this CLI invocation. @@ -37,32 +38,82 @@ def render( When `output` (a ``-o/--output`` field name) is given, the script prints that one field instead — and, as in the real command, it takes precedence over the LLM chain and the analysis sections. + + When `download_sections` (yt-dlp ``--download-sections`` specs) is given for a + downloadable URL, the generated yt-dlp call fetches only those parts of the source. """ if output is not None: llm_gateway = None # `-o` returns before the chain runs in the real command needs_download = youtube.is_downloadable_url(source) + # Sections only apply to the download path; ignore them for a local file. + sections = list(download_sections) if (needs_download and download_sections) else [] + ranges_expr, needs_re = _download_ranges(sections) parts = ( - _header_block(llm_gateway, output, needs_download=needs_download) - + _transcribe_block(merged, source, needs_download=needs_download) + _header_block( + llm_gateway, + output, + needs_download=needs_download, + needs_re=needs_re, + has_sections=ranges_expr is not None, + ) + + _transcribe_block(merged, source, needs_download=needs_download, ranges_expr=ranges_expr) + _result_block(merged, llm_gateway, output) ) parts.append("") return "\n".join(parts) +def _render_seconds(value: float) -> str: + """A Python literal for a section bound (``inf`` has no bare literal form).""" + if value == float("inf"): + return "float('inf')" + if value == float("-inf"): + return "float('-inf')" + return repr(value) + + +def _download_ranges(sections: list[str]) -> tuple[str | None, bool]: + """Render a ``download_range_func(...)`` expression for `sections`. + + Returns ``(expression_or_None, needs_re_import)`` — ``None`` when there are no + sections, and the flag is true when a chapter-regex spec means the generated + script needs ``import re``. + """ + if not sections: + return None, False + chapters, ranges, from_url = youtube.parse_download_sections(sections) + chapters_src = "[" + ", ".join(f"re.compile({c!r})" for c in chapters) + "]" + ranges_src = ( + "[" + ", ".join(f"({_render_seconds(s)}, {_render_seconds(e)})" for s, e in ranges) + "]" + ) + return f"download_range_func({chapters_src}, {ranges_src}, {from_url})", bool(chapters) + + def _header_block( - llm_gateway: dict[str, object] | None, output: str | None, *, needs_download: bool + llm_gateway: dict[str, object] | None, + output: str | None, + *, + needs_download: bool, + needs_re: bool, + has_sections: bool, ) -> list[str]: """Imports plus the api-key (and non-default environment) settings lines.""" stdlib_imports = ["import os"] if needs_download: # The download path fetches audio to a temp dir before uploading. - stdlib_imports += ["import tempfile"] + stdlib_imports.append("import tempfile") + if needs_re: + # Chapter-regex sections compile to re.compile(...) in the generated script. + stdlib_imports.append("import re") if output == "json": - stdlib_imports.insert(0, "import json") + stdlib_imports.append("import json") + stdlib_imports.sort() imports = ["import assemblyai as aai"] if needs_download: imports.append("import yt_dlp") + if has_sections: + # --download-sections builds yt-dlp's download_ranges via this helper. + imports.append("from yt_dlp.utils import download_range_func") if llm_gateway: imports.append("from openai import OpenAI") parts = [ @@ -81,7 +132,24 @@ def _header_block( return parts -def _transcribe_block(merged: dict[str, object], source: str, *, needs_download: bool) -> list[str]: +def _ydl_options_lines(ranges_expr: str | None) -> list[str]: + """The yt-dlp options dict for the download — one line, or multi-line when sections + add a ``download_ranges``/``force_keyframes_at_cuts`` pair.""" + if ranges_expr is None: + return [' {"format": "bestaudio/best", "outtmpl": f"{_tmp}/%(id)s.%(ext)s"}'] + return [ + " {", + ' "format": "bestaudio/best",', + ' "outtmpl": f"{_tmp}/%(id)s.%(ext)s",', + f' "download_ranges": {ranges_expr},', + ' "force_keyframes_at_cuts": True,', + " }", + ] + + +def _transcribe_block( + merged: dict[str, object], source: str, *, needs_download: bool, ranges_expr: str | None +) -> list[str]: """The transcriber setup, optional config, the transcribe call, and error check.""" parts = ["", "transcriber = aai.Transcriber()"] config_arg = "" @@ -97,7 +165,7 @@ def _transcribe_block(merged: dict[str, object], source: str, *, needs_download: "# AssemblyAI can't fetch this page URL itself; download the audio first.", "with tempfile.TemporaryDirectory() as _tmp:", " with yt_dlp.YoutubeDL(", - ' {"format": "bestaudio/best", "outtmpl": f"{_tmp}/%(id)s.%(ext)s"}', + *_ydl_options_lines(ranges_expr), " ) as _ydl:", f" _info = _ydl.extract_info({source!r}, download=True)", " _audio = _ydl.prepare_filename(_info)", diff --git a/aai_cli/commands/transcribe.py b/aai_cli/commands/transcribe.py index 977f6efb..163eaea3 100644 --- a/aai_cli/commands/transcribe.py +++ b/aai_cli/commands/transcribe.py @@ -29,35 +29,6 @@ app = typer.Typer() -# The PII policy strings the SDK accepts, validated client-side so a typo'd -# --redact-pii-policy fails before any upload — mirroring how an unknown --config -# key is rejected with the valid field list. -_PII_POLICY_VALUES = frozenset(policy.value for policy in aai.PIIRedactionPolicy) - - -def _validate_pii_policies(policies: list[str] | None) -> None: - unknown = [p for p in policies or [] if p not in _PII_POLICY_VALUES] - if unknown: - valid = ", ".join(sorted(_PII_POLICY_VALUES)) - raise UsageError(f"Unknown PII policy(s) {unknown}. Valid policies: {valid}.") - - -def _validate_language_flags(language_code: str | None, language_detection: bool | None) -> None: - if language_code and language_detection: - raise UsageError( - "--language-code and --language-detection can't be combined.", - suggestion="Force a language or auto-detect it, not both.", - ) - - -def _validate_speakers_expected(merged: dict[str, object]) -> None: - # Checked on the merged dict so `--config speaker_labels=true` also counts. - if merged.get("speakers_expected") and not merged.get("speaker_labels"): - raise UsageError( - "--speakers-expected only applies when diarization is enabled.", - suggestion="Add --speaker-labels.", - ) - @app.command( rich_help_panel=help_panels.TRANSCRIPTION, @@ -278,6 +249,14 @@ def transcribe( audio_end: int | None = typer.Option( None, "--audio-end", help="End offset in ms.", rich_help_panel=help_panels.OPT_CUSTOMIZATION ), + download_sections: list[str] | None = typer.Option( + None, + "--download-sections", + help="For a YouTube/podcast URL, download only part of the source (yt-dlp " + '"--download-sections" syntax, e.g. "*0:00-5:00" for the first five minutes; ' + "repeatable).", + rich_help_panel=help_panels.OPT_CUSTOMIZATION, + ), # webhooks webhook_url: str | None = typer.Option( None, @@ -370,9 +349,11 @@ def transcribe( """ def body(state: AppState, json_mode: bool) -> None: - _validate_language_flags(language_code, language_detection) + transcribe_exec.validate_language_flags( + language_code, language_detection=language_detection + ) pii_policies = config_builder.split_csv(redact_pii_policy) - _validate_pii_policies(pii_policies) + transcribe_exec.validate_pii_policies(pii_policies) flags: dict[str, object] = { "speech_model": config_builder.enum_value(speech_model), "language_code": language_code, @@ -428,7 +409,7 @@ def body(state: AppState, json_mode: bool) -> None: flags=flags, overrides=config_kv, config_file=config_file ) - _validate_speakers_expected(merged) + transcribe_exec.validate_speakers_expected(merged) if show_code: # Print-only: build the equivalent script and exit without transcribing or @@ -441,7 +422,13 @@ def body(state: AppState, json_mode: bool) -> None: ) gateway = code_gen.gateway_options(list(llm_prompt or []), model, max_tokens) output.print_code( - render_transcribe_code(merged, audio, llm_gateway=gateway, output=output_field) + render_transcribe_code( + merged, + audio, + llm_gateway=gateway, + output=output_field, + download_sections=list(download_sections or []), + ) ) return @@ -453,7 +440,11 @@ def body(state: AppState, json_mode: bool) -> None: api_key = config.resolve_api_key(profile=state.profile) with output.status("Transcribing…", json_mode=json_mode, quiet=state.quiet): transcript = transcribe_exec.run_transcription( - api_key, source, sample=sample, transcription_config=tc + api_key, + source, + sample=sample, + transcription_config=tc, + download_sections=list(download_sections or []), ) if out is not None: diff --git a/aai_cli/transcribe_exec.py b/aai_cli/transcribe_exec.py index 9a0bf859..da0c771e 100644 --- a/aai_cli/transcribe_exec.py +++ b/aai_cli/transcribe_exec.py @@ -17,6 +17,35 @@ from aai_cli import choices, client, stdio, youtube from aai_cli.errors import UsageError +# The PII policy strings the SDK accepts, validated client-side so a typo'd +# --redact-pii-policy fails before any upload — mirroring how an unknown --config +# key is rejected with the valid field list. +PII_POLICY_VALUES = frozenset(policy.value for policy in aai.PIIRedactionPolicy) + + +def validate_pii_policies(policies: list[str] | None) -> None: + unknown = [p for p in policies or [] if p not in PII_POLICY_VALUES] + if unknown: + valid = ", ".join(sorted(PII_POLICY_VALUES)) + raise UsageError(f"Unknown PII policy(s) {unknown}. Valid policies: {valid}.") + + +def validate_language_flags(language_code: str | None, *, language_detection: bool | None) -> None: + if language_code and language_detection: + raise UsageError( + "--language-code and --language-detection can't be combined.", + suggestion="Force a language or auto-detect it, not both.", + ) + + +def validate_speakers_expected(merged: dict[str, object]) -> None: + # Checked on the merged dict so `--config speaker_labels=true` also counts. + if merged.get("speakers_expected") and not merged.get("speaker_labels"): + raise UsageError( + "--speakers-expected only applies when diarization is enabled.", + suggestion="Add --speaker-labels.", + ) + def render_transform_steps(d: dict[str, Any]) -> str: """Human view of chained LLM-Gateway steps: the lone output, or each step labeled.""" @@ -56,6 +85,7 @@ def run_transcription( *, sample: bool, transcription_config: aai.TranscriptionConfig, + download_sections: list[str] | None = None, ) -> aai.Transcript: if source == "-": # Audio piped on stdin (e.g. `ffmpeg -i v.mp4 -f wav - | aai transcribe -`). @@ -72,6 +102,6 @@ def run_transcription( if youtube.is_downloadable_url(audio): # Fetch first; AssemblyAI can't read a YouTube/podcast page URL itself. with tempfile.TemporaryDirectory(prefix="aai-yt-") as td: - local = youtube.download_audio(audio, Path(td)) + local = youtube.download_audio(audio, Path(td), download_sections=download_sections) return client.transcribe(api_key, str(local), config=transcription_config) return client.transcribe(api_key, audio, config=transcription_config) diff --git a/aai_cli/youtube.py b/aai_cli/youtube.py index f1ace3b3..9b05e4dd 100644 --- a/aai_cli/youtube.py +++ b/aai_cli/youtube.py @@ -10,7 +10,7 @@ import re from pathlib import Path -from aai_cli.errors import CLIError +from aai_cli.errors import CLIError, UsageError # youtube.com/watch, youtu.be/, music.youtube.com, shorts, with or without scheme. _YOUTUBE_RE = re.compile( @@ -18,6 +18,17 @@ re.IGNORECASE, ) +# yt-dlp's own --download-sections timestamp grammar (verbatim): an optional signed +# start, a "-" separator, and an optional signed end. Both sides may be omitted +# (defaulting to 0 / inf); a leading "-" negates a bound (offset from the end). +_SECTION_RANGE_RE = re.compile( + r"""(?x)(?: + (?P-?)(?P[^-]+) + )?\s*-\s*(?: + (?P-?)(?P[^-]+) + )?""" +) + # yt-dlp's default logger prints its own "ERROR: …" line straight to stderr before the # CLI can raise its one clean error, duplicating the message. Route yt-dlp's output to # a swallow-everything logger (NullHandler, no propagation) instead. @@ -61,11 +72,69 @@ def _ytdlp_extractor_claims(url: str) -> bool: return any(ie.suitable(url) and ie.ie_key() != "Generic" for ie in gen_extractor_classes()) -def download_audio(url: str, dest_dir: Path) -> Path: +def _section_timestamp(value: str, sign: str | None) -> float: + """Seconds for one side of a ``--download-sections`` range; a "-" sign negates it.""" + from yt_dlp.utils import parse_duration + + seconds = float("inf") if value in ("inf", "infinite") else parse_duration(value) + # parse_duration returns None for an unparseable timestamp, but yt-dlp leaves it + # untyped and pyright over-narrows the inferred return to ``float``. + if seconds is None: # pyright: ignore[reportUnnecessaryComparison] + raise UsageError( + f'Invalid --download-sections time "{value}".', + suggestion='Use the form "*start-end", e.g. "*0:00-5:00".', + ) + return -seconds if sign else seconds + + +def _section_range(text: str) -> tuple[float, float]: + """Parse one ``*``-stripped ``start-end`` range into (start, end) seconds.""" + match = _SECTION_RANGE_RE.fullmatch(text) if text != "-" else None + if match is None: + raise UsageError( + f'Invalid --download-sections time range "{text}".', + suggestion='Use the form "*start-end", e.g. "*0:00-5:00".', + ) + start = _section_timestamp(match.group("start") or "0", match.group("start_sign")) + end = _section_timestamp(match.group("end") or "inf", match.group("end_sign")) + if end == float("-inf"): + raise UsageError('"-inf" is not a valid --download-sections end.') + return start, end + + +def parse_download_sections( + sections: list[str], +) -> tuple[list[str], list[tuple[float, float]], bool]: + """Split yt-dlp ``--download-sections`` specs into (chapter regexes, ranges, from-url). + + Mirrors yt-dlp's own grammar: ``*from-url`` keeps the source's own chapter range, a + ``*``-prefixed spec is one or more comma-separated ``start-end`` timestamp ranges, and + anything else is a chapter-title regex. Raises ``UsageError`` on a malformed spec. + """ + chapters: list[str] = [] + ranges: list[tuple[float, float]] = [] + from_url = False + for spec in sections: + if spec == "*from-url": + from_url = True + elif spec.startswith("*"): + ranges.extend(_section_range(chunk.strip()) for chunk in spec[1:].split(",")) + else: + try: + re.compile(spec) + except re.error as err: + raise UsageError(f'Invalid --download-sections regex "{spec}": {err}.') from err + chapters.append(spec) + return chapters, ranges, from_url + + +def download_audio(url: str, dest_dir: Path, *, download_sections: list[str] | None = None) -> Path: """Download the best audio track of `url` into `dest_dir` and return its path. Uses yt-dlp; the resulting container (m4a/webm/…) is decodable by ffmpeg - (streaming) and uploadable for transcription. + (streaming) and uploadable for transcription. `download_sections` accepts yt-dlp's + ``--download-sections`` specs (e.g. ``*0:00-5:00`` for the first five minutes), each + fetching only that part of the source. """ try: import yt_dlp @@ -77,7 +146,7 @@ def download_audio(url: str, dest_dir: Path) -> Path: suggestion="Install it: pip install yt-dlp", ) from exc - options = { + options: dict[str, object] = { "format": "bestaudio/best", "outtmpl": str(dest_dir / "%(id)s.%(ext)s"), "quiet": True, @@ -85,6 +154,19 @@ def download_audio(url: str, dest_dir: Path) -> Path: "noprogress": True, "logger": _YTDLP_LOGGER, } + if download_sections: + from yt_dlp.utils import download_range_func + + chapters, ranges, from_url = parse_download_sections(download_sections) + # yt-dlp types ``ranges`` as int tuples, but its own code (and ours) uses float + # seconds; the dict value is loosely typed for YoutubeDL either way. + options["download_ranges"] = download_range_func( + [re.compile(chapter) for chapter in chapters], + ranges, # pyright: ignore[reportArgumentType] + from_url, + ) + # Cut at exact timestamps rather than the nearest keyframe (yt-dlp's default). + options["force_keyframes_at_cuts"] = True try: # yt-dlp types `params` as a private `_Params` TypedDict, but a plain options # dict is the documented public API; pyright can't reconcile the two. diff --git a/scripts/generated_code_compile_gate.py b/scripts/generated_code_compile_gate.py index 862c7140..7ffde0c5 100644 --- a/scripts/generated_code_compile_gate.py +++ b/scripts/generated_code_compile_gate.py @@ -62,6 +62,16 @@ def main() -> int: "--show-code", ), ), + ( + "transcribe-youtube-download-sections", + ( + "transcribe", + "https://youtu.be/dtp6b76pMak", + "--download-sections", + "*0:00-5:00", + "--show-code", + ), + ), ("stream-basic", ("stream", "--show-code")), ( "stream-config-llm", diff --git a/tests/__snapshots__/test_cli_output_snapshots.ambr b/tests/__snapshots__/test_cli_output_snapshots.ambr index 021da810..25963651 100644 --- a/tests/__snapshots__/test_cli_output_snapshots.ambr +++ b/tests/__snapshots__/test_cli_output_snapshots.ambr @@ -938,6 +938,14 @@ │ spellings. │ │ --audio-start INTEGER RANGE [x>=0] Start offset in ms. │ │ --audio-end INTEGER End offset in ms. │ + │ --download-sections TEXT For a YouTube/podcast │ + │ URL, download only part │ + │ of the source (yt-dlp │ + │ "--download-sections" │ + │ syntax, e.g. │ + │ "*0:00-5:00" for the │ + │ first five minutes; │ + │ repeatable). │ ╰──────────────────────────────────────────────────────────────────────────────╯ ╭─ Webhooks ───────────────────────────────────────────────────────────────────╮ │ --webhook-url TEXT Webhook URL for completion. │ diff --git a/tests/test_code_gen.py b/tests/test_code_gen.py index f202efd5..fce952c4 100644 --- a/tests/test_code_gen.py +++ b/tests/test_code_gen.py @@ -137,6 +137,42 @@ def test_transcribe_render_youtube_passes_config_to_local_upload(): assert "transcriber.transcribe(_audio, config=config)" in code +def test_transcribe_render_download_sections_timestamp_range(): + # --download-sections renders yt-dlp's download_ranges; infinities have no bare + # literal form, so they render as float('inf')/float('-inf'). A timestamp-only spec + # needs no `import re`, and force_keyframes_at_cuts pins the cut to exact times. + code = code_gen.transcribe( + {}, + source="https://youtu.be/abc123", + download_sections=["*0:00-5:00", "*10:00-inf", "*-inf-1:00"], + ) + ast.parse(code) + assert "from yt_dlp.utils import download_range_func" in code + assert "(0.0, 300.0)" in code + assert "(600.0, float('inf'))" in code + assert "(float('-inf'), 60.0)" in code + assert '"force_keyframes_at_cuts": True,' in code + assert "import re" not in code + + +def test_transcribe_render_download_sections_chapter_imports_re(): + # A chapter-regex spec compiles to re.compile(...), so the script imports re. + code = code_gen.transcribe({}, source="https://youtu.be/abc123", download_sections=["intro"]) + ast.parse(code) + assert "import re" in code + assert "download_range_func([re.compile('intro')], [], False)" in code + + +def test_transcribe_render_download_sections_ignored_for_local_file(): + # Sections only apply to the URL download path; a local source generates no yt-dlp code. + code = code_gen.transcribe({}, source="call.mp3", download_sections=["*0:00-5:00"]) + ast.parse(code) + assert "download_range_func" not in code + assert "yt_dlp" not in code + # No sections in play means no chapter regexes, so no spurious `import re` either. + assert "import re" not in code + + def test_transcribe_render_plain_url_is_not_downloaded(): # A non-YouTube http(s) URL is uploaded straight through — no yt-dlp scaffolding. code = code_gen.transcribe({}, source="https://assembly.ai/wildfires.mp3") diff --git a/tests/test_transcribe.py b/tests/test_transcribe.py index 5bb4b893..557686d3 100644 --- a/tests/test_transcribe.py +++ b/tests/test_transcribe.py @@ -284,7 +284,10 @@ def test_transcribe_youtube_url_downloads_then_transcribes(monkeypatch, mocker, _auth() fake = tmp_path / "vid.m4a" fake.write_bytes(b"x") - monkeypatch.setattr("aai_cli.transcribe_exec.youtube.download_audio", lambda url, d: fake) + monkeypatch.setattr( + "aai_cli.transcribe_exec.youtube.download_audio", + lambda url, d, *, download_sections=None: fake, + ) tx = mocker.patch( "aai_cli.commands.transcribe.client.transcribe", autospec=True, @@ -295,13 +298,41 @@ def test_transcribe_youtube_url_downloads_then_transcribes(monkeypatch, mocker, assert tx.call_args.args[1] == str(fake) # transcribed the downloaded local file +def test_transcribe_youtube_url_forwards_download_sections(monkeypatch, mocker, tmp_path): + # --download-sections must reach youtube.download_audio so only that slice is fetched. + _auth() + fake = tmp_path / "vid.m4a" + fake.write_bytes(b"x") + seen: dict[str, object] = {} + + def _capture(url, d, *, download_sections=None): + seen["sections"] = download_sections + return fake + + monkeypatch.setattr("aai_cli.transcribe_exec.youtube.download_audio", _capture) + mocker.patch( + "aai_cli.commands.transcribe.client.transcribe", + autospec=True, + return_value=_fake_transcript(mocker), + ) + result = runner.invoke( + app, + ["transcribe", "https://youtu.be/abc", "--download-sections", "*0:00-5:00", "--json"], + ) + assert result.exit_code == 0 + assert seen["sections"] == ["*0:00-5:00"] + + def test_transcribe_podcast_page_url_downloads_then_transcribes(monkeypatch, mocker, tmp_path): # A podcast episode page (claimed by a dedicated yt-dlp extractor) routes through # the same download-first path as YouTube. _auth() fake = tmp_path / "episode.m4a" fake.write_bytes(b"x") - monkeypatch.setattr("aai_cli.transcribe_exec.youtube.download_audio", lambda url, d: fake) + monkeypatch.setattr( + "aai_cli.transcribe_exec.youtube.download_audio", + lambda url, d, *, download_sections=None: fake, + ) tx = mocker.patch( "aai_cli.commands.transcribe.client.transcribe", autospec=True, @@ -317,7 +348,7 @@ def test_transcribe_direct_audio_url_passes_through_without_download(monkeypatch # The API fetches direct audio URLs itself; no yt-dlp download must happen. _auth() - def _no_download(url, d): + def _no_download(url, d, *, download_sections=None): raise AssertionError("direct audio URLs must not be downloaded") monkeypatch.setattr("aai_cli.transcribe_exec.youtube.download_audio", _no_download) @@ -349,6 +380,21 @@ def test_transcribe_show_code_prints_without_transcribing(monkeypatch): assert "your-audio-file.mp3" not in result.output +def test_transcribe_show_code_includes_download_sections(monkeypatch): + # --show-code reflects --download-sections in the generated yt-dlp download block. + def _boom(*a, **k): + raise AssertionError("must not transcribe") + + monkeypatch.setattr("aai_cli.commands.transcribe.client.transcribe", _boom) + result = runner.invoke( + app, + ["transcribe", "https://youtu.be/abc", "--download-sections", "*0:00-5:00", "--show-code"], + ) + assert result.exit_code == 0 + assert "download_range_func([], [(0.0, 300.0)], False)" in result.output + assert '"force_keyframes_at_cuts": True,' in result.output + + def test_transcribe_show_code_without_source_uses_placeholder(monkeypatch): # --show-code never transcribes, so it must not demand a source/--sample; it emits # runnable code with a clearly-marked placeholder path instead of erroring. diff --git a/tests/test_youtube.py b/tests/test_youtube.py index ad3eda75..571e0e8c 100644 --- a/tests/test_youtube.py +++ b/tests/test_youtube.py @@ -1,10 +1,11 @@ +import re import sys import types import pytest from aai_cli import youtube -from aai_cli.errors import CLIError +from aai_cli.errors import CLIError, UsageError def test_is_youtube_url_variants(): @@ -244,3 +245,111 @@ def test_missing_ytdlp_suggests_install(tmp_path, monkeypatch): youtube.download_audio("https://youtu.be/x", tmp_path) assert "yt-dlp" in exc.value.message assert "pip install yt-dlp" in (exc.value.suggestion or "") + + +def test_parse_download_sections_timestamp_ranges(): + # A "*"-prefixed spec is one or more comma-separated start-end timestamp ranges; + # an omitted/`inf` end means "to the end", and a leading "-" negates a bound. + assert youtube.parse_download_sections(["*0:00-5:00"]) == ([], [(0.0, 300.0)], False) + assert youtube.parse_download_sections(["*10:00-inf"]) == ([], [(600.0, float("inf"))], False) + assert youtube.parse_download_sections(["*1:30-"]) == ([], [(90.0, float("inf"))], False) + # Comma-separated ranges in one spec, tolerating whitespace around each token. + assert youtube.parse_download_sections(["*0:30-1:00, 2:00-3:00"]) == ( + [], + [(30.0, 60.0), (120.0, 180.0)], + False, + ) + # "infinite" is accepted as an alias for "inf". + assert youtube.parse_download_sections(["*0:00-infinite"]) == ([], [(0.0, float("inf"))], False) + # A leading "-" on a bound negates it (offset from the end) — distinguishes the sign + # branch from a no-op. + assert youtube.parse_download_sections(["*-5:00-10:00"]) == ([], [(-300.0, 600.0)], False) + + +def test_parse_download_sections_chapters_and_from_url(): + # A non-"*" spec is a chapter-title regex; "*from-url" keeps the source's own range. + assert youtube.parse_download_sections(["intro"]) == (["intro"], [], False) + assert youtube.parse_download_sections(["*from-url"]) == ([], [], True) + # Specs combine: a chapter regex plus a timestamp range plus from-url. + assert youtube.parse_download_sections(["intro", "*0:00-1:00", "*from-url"]) == ( + ["intro"], + [(0.0, 60.0)], + True, + ) + + +@pytest.mark.parametrize( + ("spec", "needle"), + [ + ("*abc-def", 'time "abc"'), # unparseable timestamp + ("*5:00", 'time range "5:00"'), # missing the "-" separator + ("*-", 'time range "-"'), # a lone "-" is not a range + ("*1:00--inf", "-inf"), # "-inf" is not a valid end + ("(", "regex"), # malformed chapter regex + ], +) +def test_parse_download_sections_rejects_malformed(spec, needle): + with pytest.raises(UsageError) as exc: + youtube.parse_download_sections([spec]) + assert needle in exc.value.message + assert exc.value.exit_code == 2 + + +def test_download_audio_with_sections_sets_download_ranges(tmp_path, monkeypatch): + # --download-sections must reach yt-dlp as download_ranges + force_keyframes_at_cuts + # (exact cuts, not the nearest keyframe). + captured = {} + + class FakeYDL: + def __init__(self, opts): + captured["opts"] = opts + + def __enter__(self): + return self + + def __exit__(self, *exc): + return False + + def extract_info(self, url, download): + (tmp_path / "x.m4a").write_bytes(b"audio") + return {"id": "x", "ext": "m4a"} + + def prepare_filename(self, info): + return str(tmp_path / "x.m4a") + + _fake_ytdlp(monkeypatch, FakeYDL) + youtube.download_audio( + "https://youtu.be/x", tmp_path, download_sections=["*0:00-5:00", "intro"] + ) + download_ranges = captured["opts"]["download_ranges"] + assert download_ranges.ranges == [(0.0, 300.0)] + # Chapter-regex specs are compiled before reaching yt-dlp. + assert download_ranges.chapters == [re.compile("intro")] + assert captured["opts"]["force_keyframes_at_cuts"] is True + + +def test_download_audio_without_sections_omits_download_ranges(tmp_path, monkeypatch): + # The default path must not set download_ranges (downloads the whole track). + captured = {} + + class FakeYDL: + def __init__(self, opts): + captured["opts"] = opts + + def __enter__(self): + return self + + def __exit__(self, *exc): + return False + + def extract_info(self, url, download): + (tmp_path / "x.m4a").write_bytes(b"audio") + return {"id": "x", "ext": "m4a"} + + def prepare_filename(self, info): + return str(tmp_path / "x.m4a") + + _fake_ytdlp(monkeypatch, FakeYDL) + youtube.download_audio("https://youtu.be/x", tmp_path) + assert "download_ranges" not in captured["opts"] + assert "force_keyframes_at_cuts" not in captured["opts"]