diff --git a/docs/dev/rest-api/endpoints.md b/docs/dev/rest-api/endpoints.md index b4e342f3..e48f049b 100644 --- a/docs/dev/rest-api/endpoints.md +++ b/docs/dev/rest-api/endpoints.md @@ -462,6 +462,8 @@ Make a download job. It requires `service` and `title_id`. Every other field is | `cdm_only` | boolean | `null` | Force CDM-only (`true`) or vault-only (`false`) key retrieval. | | `proxy` / `no_proxy` / `no_proxy_download` / `proxy_download` | string / bool / bool / string | `null` / `false` / `false` / `null` | Proxy controls. `proxy` and `proxy_download` must be full proxy URIs unless the API key has `server_proxy`. A Control D resolver, `controld://@dns.controld.com`, counts as a full proxy URI: the server runs a forwarder for it. | | `no_folder` / `no_source` / `no_mux` | boolean | `false` | Output/mux controls. | +| `sample` | boolean/string | `null` | Download a short sample and write it with a `SAMPLE` tag on the filename in a `Sample` folder inside the output folder. `true` is 60 seconds from 2 minutes in; `"0-60"` is a window in seconds. | +| `include_sample` | boolean | `false` | Download the full title, then also cut a 60 second sample from each muxed file into a `Sample` folder next to it. | | `workers` | int | `null` | Threads per track. | | `adaptive_workers` | boolean | `false` | Scale segment workers to measured CDN throughput, up to the `workers` cap. | | `download_processes` | int | `1` | Split a track's segments across this many processes. Only engages for large batches. | diff --git a/docs/guide/cli-reference.md b/docs/guide/cli-reference.md index d65655ff..436a1cfb 100644 --- a/docs/guide/cli-reference.md +++ b/docs/guide/cli-reference.md @@ -153,6 +153,8 @@ Keep only certain track types, or skip certain track types. | `--no-folder` | Disable folder creation for TV shows. | | `--no-source` | Remove the source tag from the filename/path. | | `--no-mux` | Do not mux; keep individual track files. | +| `--sample` | Download a short sample and write it with a `SAMPLE` tag on the filename in a `Sample` folder inside the output folder. Bare `--sample` = 60 seconds from 2 minutes in; `--sample 0-60` = a custom window in seconds. | +| `--include-sample` | Download the full title, then also cut a 60 second sample from each muxed file into a `Sample` folder next to it. | | `--tag` | Group tag override. | | `--repack` | Add a `REPACK` tag to the filename. | | `--postscript` | Run a command after each output file, with `{variable}` placeholders substituted. Repeatable. Replaces the `post_scripts` config for this run. See [Post-scripts](../reference/configuration/post-scripts.md). | diff --git a/docs/guide/downloading.md b/docs/guide/downloading.md index e710ee02..792ea176 100644 --- a/docs/guide/downloading.md +++ b/docs/guide/downloading.md @@ -711,9 +711,23 @@ By default, `mkvmerge` muxes the tracks of a movie or an episode into a single M | Flag | Behaviour | Default source | | --- | --- | --- | | `--no-mux` | Do not mux; keep the individual track files. | - | +| `--sample` | Download a short sample: 60 seconds from 2 minutes in, or `--sample START-END` for a window in seconds. | - | +| `--include-sample` | Download the full title, then also cut a 60 second sample from each muxed file into a `Sample` folder. | - | | `--split-audio` | Write a separate output file per audio codec instead of merging all audio. | config `muxing.merge_audio` (on) | | `--merge-video` | Mux video tracks that share a height, range, and codec into one file, so only language varies inside a file. | config `muxing.merge_video` (off) | +!!! note "`--sample` needs a segmented track" + `--sample` downloads only the segments that hold the window, and `mkvmerge` cuts the output + to the window. The cut starts at the first video keyframe in the window, so a sample can be + some seconds short. unshackle downloads each subtitle track in full and then keeps only the + cues of the window. + + The run stops before the download starts in two conditions: a service gives a video or + audio track as one file with no segments and no segment index, or the title ends before the + window starts. + + A sidecar subtitle file with `sidecar_format: original` stays whole. + !!! note "Music is never muxed" A song skips the muxer. unshackle keeps the container that the service delivered, such as `.flac`, `.m4a`, or `.mp3`, and writes the metadata into that file. diff --git a/docs/guide/output-and-naming.md b/docs/guide/output-and-naming.md index fbdb2cc9..daed2a37 100644 --- a/docs/guide/output-and-naming.md +++ b/docs/guide/output-and-naming.md @@ -330,6 +330,22 @@ A few muxing behaviors are automatic and not configurable: The `--no-mux` flag skips muxing entirely and writes the individual track files, each with a track-type suffix, into the same output/folder structure. +## Sample files + +[`--sample`](cli-reference.md#output-muxing-files) and `--include-sample` write a sample into a +`Sample` folder inside the output folder of the title. The sample keeps the filename of the +title, with a `SAMPLE` tag at the end. The tag uses the separator of the output template. + +```text title="A movie with --include-sample" +Some.Movie.2026.1080p.EXAMPLE.WEB-DL.H.264-TAG/ +├── Some.Movie.2026.1080p.EXAMPLE.WEB-DL.H.264-TAG.mkv +└── Sample/ + └── Some.Movie.2026.1080p.EXAMPLE.WEB-DL.H.264-TAG.SAMPLE.mkv +``` + +With `--sample`, the `Sample` folder holds the only output file. If the run makes no title +folder, the `Sample` folder is in the downloads directory. + ## Tags & group naming The release-group tag is the `-TAG` portion at the end of scene-style names, which the `{tag}` variable gives. diff --git a/tests/test_sample.py b/tests/test_sample.py new file mode 100644 index 00000000..520b2f46 --- /dev/null +++ b/tests/test_sample.py @@ -0,0 +1,456 @@ +"""A ``--sample`` run downloads only the segments of its window and cuts the mux to the window.""" + +from __future__ import annotations + +import functools +import http.server +import struct +import threading +from pathlib import Path +from types import SimpleNamespace +from typing import Optional +from unittest.mock import MagicMock + +import click +import pytest +from requests import Session + +from unshackle.commands.dl import sample_name +from unshackle.core import binaries +from unshackle.core.api.handlers import validate_download_parameters +from unshackle.core.config import config +from unshackle.core.manifests.dash import DASH +from unshackle.core.services import Services +from unshackle.core.titles import Episode, Movie +from unshackle.core.tracks import Audio, Chapter, Subtitle, Tracks, Video +from unshackle.core.tracks.track import parse_sidx +from unshackle.core.utils.click_types import SAMPLE_RANGE +from unshackle.core.utils.xml import load_xml + +STARTS = [1000 + i * 4000 for i in range(100)] # 4 s segments at timescale 1000 + + +def make_video() -> Video: + return Video(id_="v", url="https://x/v", codec=Video.Codec.AVC, language="en", width=1920, height=1080, fps=24) + + +def make_audio(url: str = "https://x/a") -> Audio: + return Audio(id_="a", url=url, codec=Audio.Codec.AAC, language="en", bitrate=128000) + + +def test_no_window_keeps_every_segment() -> None: + assert make_audio().sample_slice(STARTS, 1000) == slice(None) + + +def test_the_window_is_measured_from_the_first_segment() -> None: + track = make_audio() + assert track.sample_slice(STARTS, 1000, (62, 120)) == slice(15, 30) + assert track.sample_start == 60 + + +def test_a_video_keeps_one_segment_after_the_window() -> None: + """mkvmerge ends the cut only at a video keyframe, so the video must reach past the window.""" + assert make_video().sample_slice(STARTS, 1000, (0, 60)) == slice(0, 16) + assert make_audio().sample_slice(STARTS, 1000, (0, 60)) == slice(0, 15) + + +def test_a_subtitle_keeps_every_segment() -> None: + subtitle = Subtitle(id_="s", url="https://x/s", codec=Subtitle.Codec.WebVTT, language="en") + assert subtitle.sample_slice(STARTS, 1000, (62, 120)) == slice(None) + assert subtitle.sample_start == 0 + + +def test_a_window_after_the_end_stops_the_run() -> None: + with pytest.raises(click.ClickException, match="ends before the sample window"): + make_audio().sample_slice(STARTS[:10], 1000, (120, 180)) + + +TIMELINE_MPD = """ + + + + + + + + + + +""" + +DURATION_MPD = """ + + + + + + + + +""" + +LIST_MPD = """ + + + + + https://example.com/a/ + {urls} + + + + +""".format(urls="".join(f'' for n in range(100))) + +BASE_MPD = """ + + + + https://example.com/a.mp4 + + + +""" + + +WINDOW = (62, 120) + + +def dash_names(mpd: str, window: Optional[tuple[int, int]]) -> tuple[list[str], float]: + tree = load_xml(mpd.encode()) + period = tree.find("Period") + aset = period.find("AdaptationSet") + rep = aset.find("Representation") + track = make_audio() + _, segments, _, _, _ = DASH.get_period_segments( + period, aset, rep, tree, track, "https://example.com/main.mpd", MagicMock(), window=window + ) + return [url.rsplit("/", 1)[1].split(".")[0] for url, _ in segments], track.sample_start + + +def test_dash_timeline_keeps_the_numbers_and_times_of_the_window() -> None: + names, start = dash_names(TIMELINE_MPD, WINDOW) + assert names == [f"{5 + n}-{1000 + n * 4}" for n in range(15, 30)] + assert start == 60 + + +def test_dash_fixed_duration_keeps_the_numbers_of_the_window() -> None: + names, start = dash_names(DURATION_MPD, WINDOW) + assert names == [str(5 + n) for n in range(15, 30)] + assert start == 60 + + +def test_dash_segment_list_keeps_the_urls_of_the_window() -> None: + names, start = dash_names(LIST_MPD, WINDOW) + assert names == [str(n) for n in range(15, 30)] + assert start == 60 + + +@pytest.mark.parametrize("mpd", [TIMELINE_MPD, DURATION_MPD, LIST_MPD]) +def test_a_caller_that_passes_no_window_gets_every_segment(mpd: str, monkeypatch: pytest.MonkeyPatch) -> None: + """The bitrate probe divides by the full duration, so a sample run must not cut its segment list.""" + monkeypatch.setattr(config, "sample", WINDOW) + assert len(dash_names(mpd, None)[0]) == 100 + + +@pytest.mark.parametrize(("mpd", "timed"), [(TIMELINE_MPD, True), (LIST_MPD, True), (BASE_MPD, False)]) +def test_a_track_given_as_one_file_has_no_segment_timing(mpd: str, timed: bool) -> None: + aset = load_xml(mpd.encode()).find("Period").find("AdaptationSet") + track = SimpleNamespace(data={"dash": {"adaptation_set": aset, "representation": aset.find("Representation")}}) + assert DASH.has_segment_timing(track) is timed + + +@pytest.fixture +def served(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> str: + handler = functools.partial(http.server.SimpleHTTPRequestHandler, directory=str(tmp_path)) + handler.log_message = lambda *args: None # type: ignore[attr-defined] + server = http.server.ThreadingHTTPServer(("127.0.0.1", 0), handler) + threading.Thread(target=server.serve_forever, daemon=True).start() + (tmp_path / "temp").mkdir() + monkeypatch.setattr(config.directories, "temp", tmp_path / "temp") + yield f"http://127.0.0.1:{server.server_port}" + server.shutdown() + + +def test_an_hls_subtitle_downloads_whole_in_a_sample_run( + tmp_path: Path, served: str, monkeypatch: pytest.MonkeyPatch +) -> None: + for n in range(4): + (tmp_path / f"s{n}.vtt").write_text( + f"WEBVTT\n\n00:00:{n * 10 + 1:02}.000 --> 00:00:{n * 10 + 2:02}.000\ncue {n}\n" + ) + (tmp_path / "s.m3u8").write_text( + "#EXTM3U\n#EXT-X-TARGETDURATION:10\n" + + "".join(f"#EXTINF:10.0,\ns{n}.vtt\n" for n in range(4)) + + "#EXT-X-ENDLIST\n" + ) + monkeypatch.setattr(config, "sample", (12, 18)) + subtitle = Subtitle( + id_="s", url=f"{served}/s.m3u8", codec=Subtitle.Codec.WebVTT, language="en", descriptor=Subtitle.Descriptor.HLS + ) + + subtitle.download(session=Session(), prepare_drm=lambda *a, **k: None, progress=lambda **k: None) + + text = subtitle.path.read_text() + assert all(f"cue {n}" in text for n in range(4)) + + +def test_subtitle_trim_keeps_the_cues_of_the_window_from_zero(tmp_path: Path) -> None: + subtitle = Subtitle(id_="s", url="https://x/s", codec=Subtitle.Codec.SubRip, language="en") + subtitle.path = tmp_path / "s.srt" + subtitle.path.write_text( + "1\n00:00:05,000 --> 00:00:06,000\nbefore\n\n" + "2\n00:02:10,000 --> 00:02:12,000\ninside\n\n" + "3\n00:10:00,000 --> 00:10:01,000\nafter\n" + ) + + assert subtitle.trim(120, 180) is True + + text = subtitle.path.read_text() + assert "inside" in text and "before" not in text and "after" not in text + assert "00:00:10,000 --> 00:00:12,000" in text + assert subtitle.sample_start == 120 + + +def test_subtitle_trim_leaves_an_unreadable_file_whole(tmp_path: Path) -> None: + subtitle = Subtitle(id_="s", url="https://x/s", codec=Subtitle.Codec.fTTML, language="en") + subtitle.path = tmp_path / "s.stpp" + subtitle.path.write_bytes(b"\x00\x01not a subtitle") + + assert subtitle.trim(120, 180) is False + assert subtitle.path.read_bytes() == b"\x00\x01not a subtitle" + + +def mux_command(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> list[str]: + class Stop(Exception): + pass + + captured: list[list[str]] = [] + + def fake_popen(args: list[str], **kwargs: object) -> None: + captured.append(args) + raise Stop + + monkeypatch.setattr("unshackle.core.tracks.tracks.subprocess.Popen", fake_popen) + monkeypatch.setattr(binaries, "MKVToolNix", Path("mkvmerge")) + + video, audio = make_video(), make_audio() + trimmed = Subtitle(id_="s1", url="https://x/s1", codec=Subtitle.Codec.SubRip, language="en") + whole = Subtitle(id_="s2", url="https://x/s2", codec=Subtitle.Codec.SubRip, language="fr") + for track, name in ((video, "v.mp4"), (audio, "a.mp4"), (trimmed, "s1.srt"), (whole, "s2.srt")): + track.path = tmp_path / name + track.path.write_bytes(b"\x00" * 16) + video.sample_start, audio.sample_start, trimmed.sample_start = 120.0, 119.5, 121.0 + tracks = Tracks([video, audio, trimmed, whole]) + tracks.chapters.add(Chapter("00:02:05.000", "Scene")) + + with pytest.raises(Stop): + tracks.mux("Title", delete=False, progress=lambda **k: None, output_path=tmp_path / "out.mkv") + return captured[0] + + +def test_the_mux_shifts_each_track_and_cuts_to_the_window(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(config, "sample", (121, 181)) + command = mux_command(tmp_path, monkeypatch) + + def sync_before(name: str) -> str: + file_at = command.index(str(tmp_path / name)) + options = command[:file_at] + sync_at = len(options) - 1 - options[::-1].index("--sync") + assert all(not arg.endswith((".mp4", ".srt")) for arg in command[sync_at:file_at]), ( + "the --sync is another file's" + ) + return command[sync_at + 1] + + assert command[command.index("--split") + 1] == "parts:1.500s-61.500s" + assert sync_before("v.mp4") == "0:500" + assert "--sync" not in command[: command.index(str(tmp_path / "a.mp4"))][-12:], "the earliest track is not moved" + assert sync_before("s1.srt") == "0:1500" + assert sync_before("s2.srt") == "0:-119500" + assert command[command.index("--chapter-sync") + 1] == "-119500" + assert command.index("--chapter-sync") < command.index("--chapters") + + +def test_a_full_download_adds_no_sample_options(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(config, "sample", None) + command = mux_command(tmp_path, monkeypatch) + assert not {"--split", "--sync", "--chapter-sync"} & set(command) + + +def test_sample_range_reads_a_window() -> None: + assert SAMPLE_RANGE.convert("60-120", None, None) == (60, 120) + assert SAMPLE_RANGE.convert(True, None, None) == (120, 180) + assert SAMPLE_RANGE.convert(False, None, None) is None + for bad in ("120-60", "60-60", "abc", "-5-10", "60"): + with pytest.raises(click.BadParameter): + SAMPLE_RANGE.convert(bad, None, None) + + +@pytest.mark.parametrize("sample", [True, "60-120", [60, 120]]) +def test_the_api_accepts_a_sample_window(sample: object) -> None: + assert validate_download_parameters({"sample": sample}) is None + + +@pytest.mark.parametrize("sample", ["abc", "120-60", [1, 2, 3], ["a", "b"], 5, {"start": 1}]) +def test_the_api_refuses_a_bad_sample_window(sample: object) -> None: + assert "sample" in (validate_download_parameters({"sample": sample}) or "") + + +def test_the_api_refuses_sample_with_include_sample() -> None: + assert "include_sample" in (validate_download_parameters({"sample": True, "include_sample": True}) or "") + + +def test_sample_name_uses_the_separator_of_the_template(monkeypatch: pytest.MonkeyPatch) -> None: + separators = {"movies": ".", "series": " "} + monkeypatch.setattr(config, "get_template_separator", lambda kind="movies": separators[kind]) + movie = Movie.__new__(Movie) + episode = Episode.__new__(Episode) + assert sample_name("Some.Movie.2026.1080p-TAG", movie) == "Some.Movie.2026.1080p-TAG.SAMPLE" + assert sample_name("Some Show S01E01 1080p-TAG", episode) == "Some Show S01E01 1080p-TAG SAMPLE" + + +SIDX_MPD = """ + + + + + https://example.com/a.mp4 + + + + + +""" + + +def sidx_box(count: int = 100, size: int = 5000, duration: int = 4000, timescale: int = 1000) -> bytes: + body = struct.pack(">IIIII", 0, 1, timescale, 0, 0) + struct.pack(">HH", 0, count) + body += b"".join(struct.pack(">III", size, duration, 0x90000000) for _ in range(count)) + return struct.pack(">I", 8 + len(body)) + b"sidx" + body + + +def test_parse_sidx_gives_the_time_and_byte_range_of_each_subsegment() -> None: + box = sidx_box(count=3) + end = 800 + len(box) + assert parse_sidx(box, 800) == ( + 1000, + [(0, end, end + 4999), (4000, end + 5000, end + 9999), (8000, end + 10000, end + 14999)], + ) + + +def test_parse_sidx_refuses_what_it_cannot_cut() -> None: + index_of_indexes = bytearray(sidx_box(count=2)) + index_of_indexes[32] |= 0x80 + assert parse_sidx(bytes(index_of_indexes), 0) == (0, []) + assert parse_sidx(b"\x00\x00\x00\x20moof" + b"\x00" * 24, 0) == (0, []) + assert parse_sidx(sidx_box(count=3)[:-4], 0) == (0, []) + + +def test_a_single_file_dash_track_is_cut_by_its_segment_index(monkeypatch: pytest.MonkeyPatch) -> None: + box = sidx_box() + requested: list[str] = [] + + def get(url: str, headers: dict[str, str] | None = None) -> MagicMock: + byte_range = (headers or {}).get("Range", "") + requested.append(byte_range) + content = box if byte_range == f"bytes=800-{800 + len(box) - 1}" else b"\x00" * 800 + return MagicMock(content=content, headers={"Content-Range": "bytes 0-799/900000"}) + + tree = load_xml(SIDX_MPD.format(last=800 + len(box) - 1).encode()) + period = tree.find("Period") + aset = period.find("AdaptationSet") + rep = aset.find("Representation") + track = make_audio() + track.data["dash"] = {"adaptation_set": aset, "representation": rep} + session = MagicMock(get=get) + first_media_byte = 800 + len(box) + + _, segments, timescale, starts, _ = DASH.get_period_segments( + period, aset, rep, tree, track, "https://example.com/main.mpd", session, probe_kid=False, window=WINDOW + ) + + assert DASH.has_segment_timing(track) is True + assert timescale == 1000 and starts == [n * 4000 for n in range(15, 30)] + assert segments[0] == ( + "https://example.com/a.mp4", + f"{first_media_byte + 15 * 5000}-{first_media_byte + 16 * 5000 - 1}", + ) + assert len(segments) == 15 and track.sample_start == 60 + + requested.clear() + _, whole, _, _, _ = DASH.get_period_segments( + period, aset, rep, tree, track, "https://example.com/main.mpd", session + ) + assert len(whole) == 1 and not any("800-" in r for r in requested), "a full download does not read the index" + + +@pytest.mark.parametrize( + ("args", "window"), + [ + (["--sample", "TAG", "-t", "x", "id"], (120, 180)), + (["--sample", "60-120", "TAG", "id"], (60, 120)), + (["TAG", "id"], None), + ], +) +def test_a_bare_sample_option_does_not_take_the_service_tag(args: list[str], window: object) -> None: + seen: dict[str, object] = {} + + class Group(Services): + def list_commands(self, ctx: click.Context) -> list[str]: + return ["TAG"] + + def get_command(self, ctx: click.Context, name: str) -> click.Command: + return click.Command( + name, params=[click.Option(["-t"]), click.Argument(["title"])], callback=lambda **k: None + ) + + @click.group(cls=Group) + @click.option("--sample", type=SAMPLE_RANGE, default=None) + def group(sample: object) -> None: + seen["sample"] = sample + + group(args, standalone_mode=False) + assert seen["sample"] == window + + +class RangeSession: + """Serves byte ranges of one in-memory file and records each range it was asked for.""" + + def __init__(self, data: bytes) -> None: + self.data = data + self.ranges: list[str] = [] + + def get(self, url: str, headers: dict[str, str]) -> SimpleNamespace: + self.ranges.append(headers["Range"]) + first, last = (int(n) for n in headers["Range"].removeprefix("bytes=").split("-")) + return SimpleNamespace(ok=True, content=self.data[first : last + 1]) + + +def one_file(moov_size: int, index: bytes) -> bytes: + ftyp = struct.pack(">I4s", 16, b"ftyp") + b"isom" + b"\x00" * 4 + moov = struct.pack(">I4s", moov_size, b"moov") + b"\x00" * (moov_size - 8) + return ftyp + moov + index + b"\x00" * 600_000 + + +@pytest.mark.parametrize("moov_size", [700, 200_000]) +def test_a_one_file_track_is_cut_by_the_index_in_the_file(moov_size: int, monkeypatch: pytest.MonkeyPatch) -> None: + """A large moov puts the index after the first probe, so the probe must read on.""" + monkeypatch.setattr(config, "sample", WINDOW) + box = sidx_box() + session = RangeSession(one_file(moov_size, box)) + audio = make_audio("https://x/a.mp4") + + init, byte_range = audio.sample_byte_range(session) + + media = 16 + moov_size + len(box) + assert init == session.data[: 16 + moov_size], "the init data stops before the index" + assert byte_range == f"{media + 15 * 5000}-{media + 30 * 5000 - 1}" + assert audio.sample_start == 60 + assert all(r.startswith("bytes=") for r in session.ranges), "every request is a byte range" + + +def test_a_one_file_track_with_no_index_cannot_be_cut(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(config, "sample", WINDOW) + mdat = struct.pack(">I4s", 600_000, b"mdat") + b"\x00" * (600_000 - 8) + audio = make_audio("https://x/a.mp4") + assert audio.sample_byte_range(RangeSession(one_file(700, b"")[:716] + mdat)) is None diff --git a/unshackle/commands/dl.py b/unshackle/commands/dl.py index 05c21734..f4bd8986 100644 --- a/unshackle/commands/dl.py +++ b/unshackle/commands/dl.py @@ -46,6 +46,7 @@ from unshackle.core.cdm.detect import cdm_type_stub, is_playready_cdm, is_remote from unshackle.core.config import cdm_entry_names_device, config, resolve_cdm_name, resolve_decryption from unshackle.core.console import GradientPulseBarColumn, SyncLive, console, listing_panel from unshackle.core.constants import ( + DEFAULT_SAMPLE, DOWNLOAD_ALL_DRM, DOWNLOAD_CANCELLED, DOWNLOAD_LICENCE_ONLY, @@ -107,6 +108,7 @@ from unshackle.core.utils.click_types import ( AUDIO_CODEC_LIST, LANGUAGE_RANGE, QUALITY_LIST, + SAMPLE_RANGE, SEASON_RANGE, SLOW_DELAY_RANGE, SUBTITLE_CODEC, @@ -118,7 +120,7 @@ from unshackle.core.utils.collections import ci_get, merge_dict from unshackle.core.utils.post_scripts import NO_POST_SCRIPTS, build_context, dispatch, season_context from unshackle.core.utils.redact import mask_proxy from unshackle.core.utils.selector import select_multiple -from unshackle.core.utils.subprocess import ffmpeg_decodes, ffprobe +from unshackle.core.utils.subprocess import ffmpeg_decodes, ffprobe, run_step from unshackle.core.vault import Vault from unshackle.core.vaults import Vaults @@ -259,6 +261,19 @@ def latest_episode_keys(titles: Any, latest_episode: bool, latest_episodes: Opti return list(dict.fromkeys((ep.season, ep.number) for ep in titles))[-count:] +SAMPLE_DIR = "Sample" + + +def template_type(title: Any) -> str: + """Return the output template key of a title: ``series``, ``songs`` or ``movies``.""" + return "series" if isinstance(title, Episode) else "songs" if isinstance(title, Song) else "movies" + + +def sample_name(name: str, title: Any) -> str: + """Return ``name`` with a SAMPLE tag, joined with the separator of the output template of the title.""" + return f"{name.rstrip()}{config.get_template_separator(template_type(title))}SAMPLE" + + def apply_enrich(title: Title_T, name: Optional[str], year: Optional[int], language: Optional[Language]) -> None: """Set the --enrich name, year and original language on a title. A song takes only the language. @@ -324,6 +339,41 @@ def remux_music_tracks(title: Any, log: logging.Logger) -> None: log.warning(f"Could not remux audio track {track.id} for music: {e}") +def cut_music_sample(title: Any, window: tuple[int, int], log: logging.Logger) -> None: + """Cut each audio track of a Song to the sample window. A Song has no mux step that cuts it.""" + if not isinstance(title, Song) or not binaries.FFMPEG: + return + for track in title.tracks.audio: + if not track.path or not track.path.exists(): + continue + cut_path = track.path.with_stem(f"{track.path.stem}_cut") + try: + run_step( + [ + binaries.FFMPEG, + "-v", + "error", + "-y", + "-ss", + f"{max(window[0] - track.sample_start, 0):.3f}", + "-i", + track.path, + "-t", + str(window[1] - window[0]), + "-map", + "0", + "-c", + "copy", + cut_path, + ], + output=cut_path, + label="music sample cut", + ) + os.replace(cut_path, track.path) + except RuntimeError as e: + log.warning(f"Could not cut audio track {track.id} to the sample window, so it holds whole segments: {e}") + + def apply_service_dl_overrides(ctx: click.Context, service_dl_config: dict[str, Any], log: logging.Logger) -> None: """Apply ``services..dl`` config onto ``ctx.params``. Explicit CLI/env values win. The config replaces defaults and global ``dl:`` default_map values.""" @@ -972,6 +1022,19 @@ class dl: "--no-source", is_flag=True, default=False, help="Disable the source tag from the output file name and path." ) @click.option("--no-mux", is_flag=True, default=False, help="Do not mux tracks into a container file.") + @click.option( + "--sample", + type=SAMPLE_RANGE, + default=None, + help="Download a short sample and write it to a Sample folder. " + "Use --sample for 60 seconds from 2 minutes in, or --sample START-END (e.g., --sample 0-60) for a custom window.", + ) + @click.option( + "--include-sample", + is_flag=True, + default=False, + help="Also cut a 60 second sample from each muxed file into a Sample folder next to it.", + ) @click.option( "--postscript", "postscript", @@ -1566,6 +1629,8 @@ class dl: no_folder: bool, no_source: bool, no_mux: bool, + sample: Optional[tuple[int, int]], + include_sample: bool, workers: Optional[int], adaptive_workers: bool, download_processes: int, @@ -1595,6 +1660,10 @@ class dl: postscript = NO_POST_SCRIPTS if continue_downloads: config.continue_downloads = True + config.sample = sample + if sample and include_sample: + raise click.ClickException("--sample and --include-sample cannot be used together.") + sample_dir = SAMPLE_DIR if sample else "" self.tmdb_searched = False self.search_source = None self.service_anime = bool(getattr(service, "ANIME", False)) @@ -3058,6 +3127,18 @@ class dl: self.prefer_vault_keys(title) self.cache_resolved_keys(title) + if sample and not skip_dl: + whole_files = [ + track + for track in (*title.tracks.videos, *title.tracks.audio) + if not track.can_sample(track.session or service.session) + ] + if whole_files: + raise click.ClickException( + f"--sample cannot cut {len(whole_files)} track(s) of {title}: " + f"{self.service} gives each one as one file with no segment index." + ) + dl_start_time = time.time() try: @@ -3213,6 +3294,10 @@ class dl: sub_conversions[(src, Subtitle.Codec.WebVTT.name)] = ( sub_conversions.get((src, Subtitle.Codec.WebVTT.name), 0) + 1 ) + if sample: + for subtitle in title.tracks.subtitles: + if subtitle.path and subtitle.path.exists(): + subtitle.trim(*sample) for (src, dst), count in sub_conversions.items(): log_event( "subtitle_convert", @@ -3340,6 +3425,8 @@ class dl: self.log.warning(f"VUI normalization skipped for {track.id}: {e}") remux_music_tracks(title, self.log) + if sample: + cut_music_sample(title, sample, self.log) muxed_paths = [] muxed_audio_codecs: dict[Path, Optional[Audio.Codec]] = {} @@ -3522,6 +3609,11 @@ class dl: hybrid_output_path = config.directories.temp / hybrid_filename hybrid_temp_paths.append(hybrid_output_path) + if len({track.sample_start for track in resolution_tracks}) > 1: + raise click.ClickException( + "--sample cut the HDR10 and Dolby Vision tracks at different " + "points, so a hybrid of them would not line up." + ) Hybrid(resolution_tracks, self.service, output_name=hybrid_filename) if container: @@ -3579,6 +3671,8 @@ class dl: media_info = MediaInfo.parse(muxed_path) if muxed_path else None if media_info: base_filename = title.get_filename(media_info, show_service=not no_source) + if sample: + base_filename = sample_name(base_filename, title) else: base_filename = str(title) @@ -3594,6 +3688,7 @@ class dl: sidecar_dir /= title.get_filename( media_info, show_service=not no_source, folder=True ) + sidecar_dir /= sample_dir sidecar_targets.append((muxed_path, base_filename, sidecar_dir)) with console.status("Saving subtitle sidecar files..."): @@ -3674,11 +3769,14 @@ class dl: media_info = MediaInfo.parse(sample_track.path) final_dir /= title.get_filename(media_info, show_service=not no_source, folder=True) + final_dir /= sample_dir final_dir.mkdir(parents=True, exist_ok=True) for track_path in muxed_paths: media_info = MediaInfo.parse(track_path) base_filename = title.get_filename(media_info, show_service=not no_source) + if sample: + base_filename = sample_name(base_filename, title) track = next((t for t in title.tracks if t.path == track_path), None) if track: @@ -3710,6 +3808,8 @@ class dl: media_info = MediaInfo.parse(muxed_path) final_dir = self.output_dir or config.directories.downloads final_filename = title.get_filename(media_info, show_service=not no_source) + if sample: + final_filename = sample_name(final_filename, title) audio_codec_suffix = muxed_audio_codecs.get(muxed_path) if not no_folder and ( @@ -3717,13 +3817,11 @@ class dl: or (isinstance(title, Movie) and config.get_folder_template("movies")) ): final_dir /= title.get_filename(media_info, show_service=not no_source, folder=True) + final_dir /= sample_dir final_dir.mkdir(parents=True, exist_ok=True) final_path = final_dir / f"{final_filename}{muxed_path.suffix}" - template_type = ( - "series" if isinstance(title, Episode) else "songs" if isinstance(title, Song) else "movies" - ) - sep = config.get_template_separator(template_type) + sep = config.get_template_separator(template_type(title)) if final_path.exists() and audio_codec_suffix and append_audio_codec_suffix: final_filename = f"{final_filename.rstrip()}{sep}{audio_codec_suffix.name}" @@ -3754,6 +3852,8 @@ class dl: self.anime_hint(title), session=service.session, ) + if include_sample: + self.write_sample(final_path, title, media_info) post_script_context = build_context( title, @@ -3807,6 +3907,32 @@ class dl: console.print(Padding(f"Processed all titles in [progress.elapsed]{dl_time}", (0, 5, 1, 5))) + def write_sample(self, source: Path, title: Title_T, media_info: MediaInfo) -> None: + """Cut a 60 second sample from a finished file into a ``Sample`` folder next to it. + + The sample starts 2 minutes in. A file shorter than 3 minutes gives its middle 60 seconds. + """ + duration = float(media_info.general_tracks[0].duration or 0) / 1000 + start, end = DEFAULT_SAMPLE + length = end - start + if duration < end: + start = max(0.0, (duration - length) / 2) + output = source.parent / SAMPLE_DIR / f"{sample_name(source.stem, title)}{source.suffix}" + output.parent.mkdir(exist_ok=True) + returncode = subprocess.run( + [ + str(binaries.MKVToolNix), + "-q", + "-o", + str(output), + "--split", + f"parts:{start:.3f}s-{start + length:.3f}s", + source, + ] + ).returncode + if returncode > 1 or not output.exists(): + self.log.warning(f"mkvmerge wrote no sample for {source.name}") + def anime_hint(self, title: Optional[Title_T] = None) -> bool: """Whether metadata lookups for this title should prefer AniList.""" per_title = getattr(title, "anime", None) diff --git a/unshackle/core/api/download_manager.py b/unshackle/core/api/download_manager.py index 3f95c585..366b2e04 100644 --- a/unshackle/core/api/download_manager.py +++ b/unshackle/core/api/download_manager.py @@ -418,6 +418,15 @@ def perform_download( if params.get("export"): params["export"] = bool(params["export"]) + sample_raw = params.get("sample") + if sample_raw is not None: + from unshackle.core.utils.click_types import SAMPLE_RANGE + + try: + params["sample"] = SAMPLE_RANGE.convert(sample_raw, None, None) + except click.BadParameter as exc: + raise Exception(f"Invalid sample parameter: {exc}") + slow_raw = params.get("slow") if slow_raw is not None and not isinstance(slow_raw, tuple): if isinstance(slow_raw, bool): @@ -633,6 +642,8 @@ def perform_download( no_folder=params.get("no_folder", False), no_source=params.get("no_source", False), no_mux=params.get("no_mux", False), + sample=params.get("sample"), + include_sample=params.get("include_sample", False), workers=params.get("workers"), adaptive_workers=params.get("adaptive_workers", False), download_processes=params.get("download_processes", 1), diff --git a/unshackle/core/api/handlers.py b/unshackle/core/api/handlers.py index 8dc0d6d8..a96e1593 100644 --- a/unshackle/core/api/handlers.py +++ b/unshackle/core/api/handlers.py @@ -92,6 +92,8 @@ DEFAULT_DOWNLOAD_PARAMS = { "no_folder": False, "no_source": False, "no_mux": False, + "sample": None, + "include_sample": False, "workers": None, "adaptive_workers": False, "download_processes": 1, @@ -1831,6 +1833,16 @@ def validate_download_parameters(data: Dict[str, Any]) -> Optional[str]: except APIError as e: return e.message + if data.get("sample") is not None: + from unshackle.core.utils.click_types import SAMPLE_RANGE + + try: + SAMPLE_RANGE.convert(data["sample"], None, None) + except click.BadParameter: + return "sample must be true or a 'START-END' window in seconds, such as '60-120'." + if data.get("sample") and data.get("include_sample"): + return "sample and include_sample cannot be used together." + # The serve user would create and write any directory a client names; keep it under downloads. if data.get("output_dir"): root = config.directories.downloads.resolve() diff --git a/unshackle/core/api/routes.py b/unshackle/core/api/routes.py index d011b08a..dee6ea2a 100644 --- a/unshackle/core/api/routes.py +++ b/unshackle/core/api/routes.py @@ -1005,6 +1005,14 @@ async def download(request: web.Request) -> web.Response: no_mux: type: boolean description: Do not mux tracks into a container file (default - false) + include_sample: + type: boolean + description: Also cut a 60 second sample from each muxed file into a Sample folder next to it (default - false) + sample: + oneOf: + - type: boolean + - type: string + description: Download a short sample and write it to a Sample folder. `true` for 60 seconds from 2 minutes in, or a `"START-END"` string in seconds (e.g., `"0-60"`) (default - null) workers: type: integer description: Max workers/threads per track download (default - None) diff --git a/unshackle/core/config.py b/unshackle/core/config.py index 994b1d26..9ddcef4d 100644 --- a/unshackle/core/config.py +++ b/unshackle/core/config.py @@ -193,6 +193,7 @@ class Config: # mask local base dirs (install root/venv/home) in logged paths; False shows full paths self.redact_paths: bool = kwargs.get("redact_paths", True) self.continue_downloads: bool = kwargs.get("continue_downloads", False) + self.sample: Optional[tuple[int, int]] = None self.post_scripts: list = kwargs.get("post_scripts") or [] self.language_tags: dict = kwargs.get("language_tags") or {} diff --git a/unshackle/core/constants.py b/unshackle/core/constants.py index ffd84427..5118cf0b 100644 --- a/unshackle/core/constants.py +++ b/unshackle/core/constants.py @@ -3,6 +3,8 @@ from typing import TypeVar, Union DOWNLOAD_CANCELLED = Event() DOWNLOAD_LICENCE_ONLY = Event() + +DEFAULT_SAMPLE = (120, 180) DOWNLOAD_ALL_DRM = Event() diff --git a/unshackle/core/manifests/dash.py b/unshackle/core/manifests/dash.py index 4e5b37af..edf8b395 100644 --- a/unshackle/core/manifests/dash.py +++ b/unshackle/core/manifests/dash.py @@ -34,7 +34,11 @@ from unshackle.core.drm.verify import decrypt_track from unshackle.core.events import events from unshackle.core.session import RnetSession from unshackle.core.tracks import Audio, DownloadContext, Subtitle, Tracks, Video, resume -from unshackle.core.tracks.track import DRM_PREFERENCE_TYPES, assert_fragments_decrypted +from unshackle.core.tracks.track import ( + DRM_PREFERENCE_TYPES, + assert_fragments_decrypted, + parse_sidx, +) from unshackle.core.utilities import is_close_match, log_event, try_ensure_utf8 from unshackle.core.utils.redact import safe_display_url from unshackle.core.utils.xml import load_xml @@ -458,6 +462,7 @@ class DASH: track_url=track.url, session=session, probe_kid=period_idx == 0, + window=config.sample, ) if period_idx == 0: @@ -471,12 +476,20 @@ class DASH: seen_segments.add(seg) segments.append(seg) segment_durations.extend(p_durations) + # ponytail: a sample comes from the first content period only; map the window across + # periods if a manifest with a short first period needs it + if config.sample and segments: + break if not segments: log.error("Could not find a way to get segments from this MPD manifest.") log.debug(track.url) sys.exit(1) + if config.sample and not segment_durations and not isinstance(track, Subtitle): + log.error("--sample cannot cut this track: the manifest gives it as one file with no segment index.") + sys.exit(1) + # TODO: Should we floor/ceil/round, or is int() ok? track.data["dash"]["timescale"] = int(segment_timescale) track.data["dash"]["segment_durations"] = segment_durations @@ -533,7 +546,7 @@ class DASH: # When every segment is a byte range of one parent resource and init is its [0, len) # prefix, the whole resource is itself a valid MP4, so one direct download to save_path # replaces the per-segment files, init write and merge pass. - collapse_single_url = DASH.collapsible_single_url( + collapse_single_url = not config.sample and DASH.collapsible_single_url( isinstance(track, Subtitle), segments, len(init_data) if init_data is not None else None, @@ -753,6 +766,26 @@ class DASH: progress(downloaded="Downloaded") + @staticmethod + def find_segment_element(name: str, adaptation_set: Element, representation: Element) -> Optional[Element]: + """Find a segment description element on the representation, or else on its adaptation set.""" + element = representation.find(name) + return element if element is not None else adaptation_set.find(name) + + @staticmethod + def has_segment_timing(track: AnyTrack) -> bool: + """Return False when the manifest gives the track as one file with no index, so a sample cannot cut it.""" + dash = track.data.get("dash") or {} + representation, adaptation_set = dash.get("representation"), dash.get("adaptation_set") + if representation is None or adaptation_set is None: + return True + segment_base = DASH.find_segment_element("SegmentBase", adaptation_set, representation) + return ( + DASH.merge_segment_templates(adaptation_set, representation) is not None + or DASH.find_segment_element("SegmentList", adaptation_set, representation) is not None + or (segment_base is not None and bool(segment_base.get("indexRange"))) + ) + @staticmethod def collapsible_single_url( is_subtitle: bool, @@ -844,6 +877,7 @@ class DASH: track_url: str, session: Union[Session, RnetSession], probe_kid: bool = True, + window: Optional[tuple[int, int]] = None, ) -> tuple[ Optional[bytes], list[tuple[str, Optional[str]]], @@ -877,13 +911,9 @@ class DASH: segment_template = DASH.merge_segment_templates(adaptation_set, representation) - segment_list = representation.find("SegmentList") - if segment_list is None: - segment_list = adaptation_set.find("SegmentList") + segment_list = DASH.find_segment_element("SegmentList", adaptation_set, representation) - segment_base = representation.find("SegmentBase") - if segment_base is None: - segment_base = adaptation_set.find("SegmentBase") + segment_base = DASH.find_segment_element("SegmentBase", adaptation_set, representation) segments: list[tuple[str, Optional[str]]] = [] segment_timescale: float = 0 @@ -930,6 +960,10 @@ class DASH: segment_durations.append(current_time) current_time += int(s.get("d")) + keep = track.sample_slice(segment_durations, segment_timescale, window) if window else slice(None) + start_number += keep.start or 0 + segment_durations = segment_durations[keep] + if not end_number: end_number = start_number + len(segment_durations) - 1 @@ -956,7 +990,11 @@ class DASH: segment_count = math.ceil(period_duration / (segment_duration / segment_timescale)) end_number = start_number + segment_count - 1 - for s in range(start_number, end_number + 1): + numbers = range(start_number, end_number + 1) + if window: + starts = [n * segment_duration for n in range(len(numbers))] + numbers = numbers[track.sample_slice(starts, segment_timescale, window)] + for s in numbers: segments.append( ( DASH.replace_fields( @@ -994,6 +1032,10 @@ class DASH: track_kid = track.get_key_id(init_data) segment_urls = segment_list.findall("SegmentURL") + list_duration = float(segment_list.get("duration") or 0) + if window and list_duration: + starts = [n * list_duration for n in range(len(segment_urls))] + segment_urls = segment_urls[track.sample_slice(starts, segment_timescale, window)] for segment_url in segment_urls: media_url = segment_url.get("media") if not media_url: @@ -1022,7 +1064,20 @@ class DASH: if total_size: media_range = f"{len(init_data)}-{total_size}" - segments.append((rep_base_url, media_range)) + subsegments: list[tuple[int, int, int]] = [] + index_range = segment_base.get("indexRange") + if window and index_range and not isinstance(track, Subtitle): + res = session.get(url=rep_base_url, headers={"Range": f"bytes={index_range}"}) + res.raise_for_status() + segment_timescale, subsegments = parse_sidx(res.content, int(index_range.split("-")[0])) + + if subsegments: + segment_durations = [time for time, _, _ in subsegments] + keep = track.sample_slice(segment_durations, segment_timescale, window) + segment_durations = segment_durations[keep] + segments.extend((rep_base_url, f"{first}-{last}") for _, first, last in subsegments[keep]) + else: + segments.append((rep_base_url, media_range)) elif rep_base_url: segments.append((rep_base_url, None)) diff --git a/unshackle/core/manifests/hls.py b/unshackle/core/manifests/hls.py index 2d2e1ea7..bdd77e3a 100644 --- a/unshackle/core/manifests/hls.py +++ b/unshackle/core/manifests/hls.py @@ -11,6 +11,7 @@ import subprocess import sys import time from contextlib import ExitStack, closing +from itertools import accumulate from pathlib import Path from typing import Any, Optional, Union from urllib.parse import urljoin, urlparse @@ -685,6 +686,12 @@ class HLS: ) hls_data["unwanted_segments"] = [s.absolute_uri for s in unwanted_segments] + if config.sample: + content = [segment for segment in master.segments if segment not in unwanted_segments] + starts = list(accumulate((segment.duration or 0 for segment in content[:-1]), initial=0.0)) + del content[track.sample_slice(starts, window=config.sample)] + unwanted_segments += content + # Get session DRM as fallback but prefer media playlist keys for accurate KID matching if track.drm: session_drm = track.get_drm_for_cdm(cdm) @@ -796,17 +803,17 @@ class HLS: range_offset = 0 for segment in master.segments: + if segment.byterange: + byte_range = HLS.calculate_byte_range(segment.byterange, range_offset) + range_offset = int(byte_range.split("-")[1]) + 1 + else: + byte_range = None + if segment in unwanted_segments: continue segment_durations.append(int(segment.duration)) - if segment.byterange: - byte_range = HLS.calculate_byte_range(segment.byterange, range_offset) - range_offset = int(byte_range.split("-")[0]) - else: - byte_range = None - urls.append( { "url": urljoin(segment.base_uri, segment.uri), diff --git a/unshackle/core/manifests/ism.py b/unshackle/core/manifests/ism.py index a29f71ff..30bbbce7 100644 --- a/unshackle/core/manifests/ism.py +++ b/unshackle/core/manifests/ism.py @@ -259,6 +259,7 @@ class ISM: ) track_urls: list[str] = [] + track_times: list[int] = [] fragment_time = 0 fragments = stream_index.findall("c") # MS-SSTR UrlPattern; regex over str.format so {Bitrate}/{start_time} @@ -276,6 +277,7 @@ class ISM: first_time = int(fragments[0].get("t") or 0) if first_time != 0: track_urls.append(re.sub(r"\{start[ _]time\}", "0", url_template)) + track_times.append(first_time) for idx, frag in enumerate(fragments): fragment_time = int(frag.get("t", fragment_time)) @@ -292,6 +294,7 @@ class ISM: duration_frag = (next_time - fragment_time) // repeat for _ in range(repeat): track_urls.append(re.sub(r"\{start[ _]time\}", str(fragment_time), url_template)) + track_times.append(fragment_time) fragment_time += duration_frag track_id = hashlib.md5( @@ -312,6 +315,8 @@ class ISM: "stream_index": stream_index, "quality_level": ql, "segments": track_urls, + "segment_times": track_times, + "timescale": int(stream_index.get("TimeScale") or self.manifest.get("TimeScale") or 10000000), } } @@ -414,6 +419,9 @@ class ISM: session.proxies.update({"all": proxy}) segments: list[str] = track.data["ism"]["segments"] + segment_times: list[int] = track.data["ism"].get("segment_times") or [] + if segment_times: + segments = segments[track.sample_slice(segment_times, track.data["ism"]["timescale"], config.sample)] session_drm = None if track.drm: diff --git a/unshackle/core/services.py b/unshackle/core/services.py index 5212241f..c5f40d96 100644 --- a/unshackle/core/services.py +++ b/unshackle/core/services.py @@ -15,6 +15,7 @@ import click from unshackle.core import binaries from unshackle.core.config import config +from unshackle.core.constants import DEFAULT_SAMPLE from unshackle.core.service import Service from unshackle.core.service_repo import DirtyServiceRepo, head, is_repo_spec, refresh_repo, resolve_service_repo from unshackle.core.utilities import import_module_by_path @@ -360,16 +361,17 @@ class Services(click.Group): super().format_help(ctx, formatter) def parse_args(self, ctx: click.Context, args: list[str]) -> list[str]: - """Preprocess ``--slow`` so it can take an optional range value before Click parses the arguments.""" + """Preprocess ``--slow`` and ``--sample`` so each can take an optional range value before Click parses.""" + defaults = {"--slow": "60-120", "--sample": "{}-{}".format(*DEFAULT_SAMPLE)} processed = [] i = 0 while i < len(args): - if args[i] == "--slow": + if args[i] in defaults: if i + 1 < len(args) and re.match(r"^\d+-\d+$", args[i + 1]): - processed.append(f"--slow={args[i + 1]}") + processed.append(f"{args[i]}={args[i + 1]}") i += 2 else: - processed.append("--slow=60-120") + processed.append(f"{args[i]}={defaults[args[i]]}") i += 1 else: processed.append(args[i]) diff --git a/unshackle/core/tracks/subtitle.py b/unshackle/core/tracks/subtitle.py index f0e5d0ab..5d64d83c 100644 --- a/unshackle/core/tracks/subtitle.py +++ b/unshackle/core/tracks/subtitle.py @@ -9,7 +9,7 @@ from enum import Enum from functools import partial from io import BytesIO from pathlib import Path -from typing import Any, Callable, Iterable, Optional, Union +from typing import Any, Callable, Iterable, Optional, Sequence, Union import pycaption import requests @@ -1001,6 +1001,32 @@ class Subtitle(Track): return captions, language + def sample_slice( + self, starts: Sequence[float], timescale: float = 1, window: Optional[tuple[int, int]] = None + ) -> slice: + """Keep every segment: ``trim`` cuts a subtitle by its cue times after the download.""" + return slice(None) + + def trim(self, start: float, end: float) -> bool: + """Keep the cues between ``start`` and ``end`` seconds and move them to start at zero. + + Returns False and leaves the file whole when pysubs2 cannot read or write its format. + """ + import pysubs2 + + low, high = round(start * 1000), round(end * 1000) + try: + subs = pysubs2.load(str(self.path), encoding="utf-8") + subs.events = [event for event in subs.events if event.end > low and event.start < high] + for event in subs.events: + event.start = max(event.start, low) - low + event.end = min(event.end, high) - low + subs.save(str(self.path)) + except (pysubs2.exceptions.Pysubs2Error, UnicodeDecodeError): + return False + self.sample_start = start + return True + def strip_hearing_impaired(self) -> None: """ Remove the SDH captions. diff --git a/unshackle/core/tracks/track.py b/unshackle/core/tracks/track.py index fa804fcb..2fcafe0f 100644 --- a/unshackle/core/tracks/track.py +++ b/unshackle/core/tracks/track.py @@ -3,18 +3,23 @@ import html import logging import re import shutil +import struct import subprocess +from bisect import bisect_left, bisect_right from collections import defaultdict from copy import copy from dataclasses import dataclass from enum import Enum from functools import partial from pathlib import Path -from typing import Any, Callable, Iterable, Optional, Union +from typing import Any, Callable, Iterable, Optional, Sequence, Union from uuid import UUID from zlib import crc32 +import click +from construct import ConstructError from langcodes import Language +from pymp4.parser import Box from requests import Response, Session from requests.adapters import HTTPAdapter, Retry @@ -364,6 +369,31 @@ def assert_fragments_decrypted(path: Path) -> None: ) +SAMPLE_PROBE_BYTES = 64 * 1024 +SAMPLE_PROBE_LIMIT = 8 * 1024 * 1024 + + +def parse_sidx(data: bytes, offset: int) -> tuple[int, list[tuple[int, int, int]]]: + """Read a ``sidx`` box that starts at byte ``offset`` of its file. + + Returns the timescale and one (start time, first byte, last byte) for each subsegment. The list + is empty when the box is not a ``sidx`` or when it points at more indexes and not at media. + """ + try: + box = Box.parse(data) + except ConstructError: + return 0, [] + if box.type != b"sidx" or any(reference.reference_type != "MEDIA" for reference in box.references): + return 0, [] + start, byte = box.earliest_presentation_time, offset + box.end + box.first_offset + subsegments: list[tuple[int, int, int]] = [] + for reference in box.references: + subsegments.append((start, byte, byte + reference.referenced_size - 1)) + start += reference.segment_duration + byte += reference.referenced_size + return box.timescale, subsegments + + @dataclass class DownloadContext: """Shared arguments passed to each manifest's ``download_track``.""" @@ -396,6 +426,8 @@ class Track: DASH = 3 # https://en.wikipedia.org/wiki/Dynamic_Adaptive_Streaming_over_HTTP ISM = 4 # https://learn.microsoft.com/en-us/silverlight/smooth-streaming + sample_tail = 0 + def __init__( self, url: Union[str, list[str]], @@ -497,6 +529,77 @@ class Track: # TODO: Currently using OnFoo event naming, change to just segment_filter self.OnSegmentFilter: Optional[Callable] = None + self.sample_start: float = 0.0 + self.sample_bytes: Optional[tuple[bytes, str]] = None + + def sample_slice( + self, starts: Sequence[float], timescale: float = 1, window: Optional[tuple[int, int]] = None + ) -> slice: + """Return the slice of segments that hold ``window``, from the start time of each segment. + + Sets ``sample_start`` to the start time of the first kept segment, in seconds. A video keeps + ``sample_tail`` segments after the window: mkvmerge ends the cut only at a video keyframe, and + with no such keyframe the audio and the subtitles continue to their own end. + """ + self.sample_start = 0.0 + if not window or not starts: + return slice(None) + low, high = (starts[0] + seconds * timescale for seconds in window) + if low > starts[-1] and len(starts) > 1: + raise click.ClickException( + "The title ends before the sample window starts. Use --sample START-END with an earlier window." + ) + first = max(bisect_right(starts, low) - 1, 0) + self.sample_start = (starts[first] - starts[0]) / timescale + return slice(first, max(bisect_left(starts, high) + self.sample_tail, first + 1)) + + def sample_byte_range(self, session: Union[Session, "RnetSession"]) -> Optional[tuple[bytes, str]]: + """Find the bytes of a one-file track that hold the sample window. + + Reads the start of the file for a ``sidx`` box. Returns the bytes before that box, which hold + the init data, and the byte range of the window. Returns None when the file has no such index. + """ + if self.sample_bytes or not isinstance(self.url, str): + return self.sample_bytes + head = b"" + + def read_to(end: int) -> None: + nonlocal head + if end > len(head): + res = session.get(self.url, headers={"Range": f"bytes={len(head)}-{end - 1}"}) + if res.ok: + head += res.content + + pos = 0 + while pos < SAMPLE_PROBE_LIMIT: + if pos + 8 > len(head): + read_to(pos + SAMPLE_PROBE_BYTES) + if pos + 8 > len(head): + return None + size, kind = struct.unpack(">I4s", head[pos : pos + 8]) + if size < 8: + return None + if kind != b"sidx": + pos += size + continue + read_to(pos + size) + timescale, subsegments = parse_sidx(head[pos : pos + size], pos) + if not subsegments: + return None + kept = subsegments[self.sample_slice([start for start, _, _ in subsegments], timescale, config.sample)] + self.sample_bytes = head[:pos], f"{kept[0][1]}-{kept[-1][2]}" + return self.sample_bytes + return None + + def can_sample(self, session: Union[Session, "RnetSession"]) -> bool: + """Return False when a sample cannot cut this track, because it is one file with no segment index.""" + if self.descriptor == self.Descriptor.URL: + return bool(self.sample_byte_range(session)) + if self.descriptor == self.Descriptor.DASH: + from unshackle.core.manifests import DASH + + return DASH.has_segment_timing(self) + return True def __repr__(self) -> str: return "{name}({items})".format( @@ -741,8 +844,17 @@ class Track: if DOWNLOAD_LICENCE_ONLY.is_set(): progress(downloaded="[yellow]SKIPPED") else: + sample_range = None + if config.sample and track_type in ("Video", "Audio"): + sample_range = self.sample_byte_range(dl_session) + if not sample_range: + raise click.ClickException( + "--sample cannot cut this track: it is one file with no segment index." + ) for status_update in self.downloader( - urls=self.url, + urls=[{"url": self.url, "headers": {"Range": f"bytes={sample_range[1]}"}}] + if sample_range + else self.url, output_dir=save_path.parent, filename=save_path.name, headers=dl_session.headers, @@ -760,6 +872,14 @@ class Track: status_update["downloaded"] = f"URL {downloaded}" progress(**status_update) + if sample_range: + media_path = save_path.with_name(f"{save_path.name}.media") + save_path.replace(media_path) + with save_path.open("wb") as joined, media_path.open("rb") as media: + joined.write(sample_range[0]) + shutil.copyfileobj(media, joined) + media_path.unlink() + self.path = save_path events.emit(events.Types.TRACK_DOWNLOADED, track=self) diff --git a/unshackle/core/tracks/tracks.py b/unshackle/core/tracks/tracks.py index 429c5fbf..5b687d0a 100644 --- a/unshackle/core/tracks/tracks.py +++ b/unshackle/core/tracks/tracks.py @@ -34,6 +34,7 @@ from unshackle.core.utilities import ( valid_language, ) from unshackle.core.utils.collections import as_list, flatten +from unshackle.core.utils.subprocess import run_step MP4BOX_PROGRESS = re.compile(r"\((\d{1,3})/100\)") MP4BOX_OPEN_FAILED = re.compile(r"error while opening|Error opening file|Invalid IsoMedia File", re.I) @@ -495,6 +496,11 @@ class Tracks: selected.append(track) return selected + @property + def sample_base(self) -> float: + """The earliest ``sample_start`` of the video and audio tracks, which a sample mux uses as time zero.""" + return min((t.sample_start for t in (*self.videos, *self.audio)), default=0.0) + def _mux_mp4( self, title: str, @@ -524,6 +530,7 @@ class Tracks: progress = progress or partial(lambda **kwargs: None) cl = [str(mp4box)] + sample_base = self.sample_base names: list[str] = [] cleaned: list[Path] = [] track_id = 0 @@ -557,7 +564,8 @@ class Tracks: elif candidate.exists(): candidate.unlink() - cl.extend(["-add", f"{source}:lang={track.language}"]) + delay = round((track.sample_start - sample_base) * 1000) if config.sample else 0 + cl.extend(["-add", f"{source}:lang={track.language}" + (f":delay={delay}" if delay else "")]) name = "" if isinstance(track, Video) else (track.get_track_name() or "") if name: names.extend(["-name", f"{track_id}={name}"]) @@ -568,7 +576,7 @@ class Tracks: cl.extend(names) chapters_path = None - if self.chapters: + if self.chapters and not config.sample: chapters_path = config.directories.temp / config.filenames.chapters.format( title=sanitize_filename(title), random=f"{self.chapters.id}_{os.getpid()}_{threading.get_ident()}" ) @@ -632,6 +640,19 @@ class Tracks: returncode = p.wait() if returncode < 2 and any(MP4BOX_OPEN_FAILED.search(line) for line in errors): returncode = 2 + if config.sample and returncode == 0: + start, end = (seconds - sample_base for seconds in config.sample) + cut_path = output_path.with_suffix(".cut.mp4") + try: + run_step( + [mp4box, "-splitx", f"{start:.3f}:{end:.3f}", output_path, "-out", cut_path], + output=cut_path, + label="MP4Box sample cut", + ) + os.replace(cut_path, output_path) + except RuntimeError as e: + returncode = 2 + errors.append(f"MP4Box could not cut the sample window: {e}") mux_duration_ms = round((time.monotonic() - mux_start_time) * 1000, 1) output_size = output_path.stat().st_size if output_path.exists() else 0 @@ -729,6 +750,16 @@ class Tracks: "--no-date", # remove dates from the output for security ] + sample_base = self.sample_base + + def sample_sync(start: float) -> list[str]: + shift = round((start - sample_base) * 1000) + return ["--sync", f"0:{shift}"] if config.sample and shift else [] + + if config.sample: + start, end = (seconds - sample_base for seconds in config.sample) + cl.extend(["--split", f"parts:{start:.3f}s-{end:.3f}s"]) + if config.muxing.get("set_title", True): cl.extend(["--title", title]) @@ -782,6 +813,7 @@ class Tracks: f"0:{vt.is_original_lang}", "--compression", "0:none", + *sample_sync(vt.sample_start), ] # Add FPS fix if needed (typically for hybrid mode to prevent sync issues) @@ -839,6 +871,7 @@ class Tracks: f"0:{at.is_original_lang}", "--compression", "0:none", + *sample_sync(at.sample_start), "(", str(at.path), ")", @@ -872,6 +905,7 @@ class Tracks: f"0:{st.is_original_lang}", "--compression", "0:none", + *sample_sync(st.sample_start), "(", str(st.path), ")", @@ -883,6 +917,8 @@ class Tracks: title=sanitize_filename(title), random=f"{self.chapters.id}_{os.getpid()}_{threading.get_ident()}" ) self.chapters.dump(chapters_path, fallback_name=config.chapter_fallback_name) + if config.sample and sample_base: + cl.extend(["--chapter-sync", str(round(-sample_base * 1000))]) cl.extend(["--chapter-charset", "UTF-8", "--chapters", str(chapters_path)]) else: chapters_path = None diff --git a/unshackle/core/tracks/video.py b/unshackle/core/tracks/video.py index 07cad6f6..260c3016 100644 --- a/unshackle/core/tracks/video.py +++ b/unshackle/core/tracks/video.py @@ -20,6 +20,8 @@ from unshackle.core.utils.subprocess import log_tool_run class Video(Track): + sample_tail = 1 + class Codec(str, Enum): AVC = "H.264" HEVC = "H.265" diff --git a/unshackle/core/utils/click_types.py b/unshackle/core/utils/click_types.py index d1bbe655..a1e46901 100644 --- a/unshackle/core/utils/click_types.py +++ b/unshackle/core/utils/click_types.py @@ -6,6 +6,7 @@ import click from click.shell_completion import CompletionItem from pywidevine.cdm import Cdm as WidevineCdm +from unshackle.core.constants import DEFAULT_SAMPLE from unshackle.core.tracks.audio import Audio from unshackle.core.tracks.subtitle import Subtitle from unshackle.core.tracks.video import Video @@ -505,6 +506,26 @@ class SlowDelayRange(click.ParamType): return (low, high) +class SampleRange(click.ParamType): + """Parses a sample window string like '60-120' into a tuple of (start, end) seconds.""" + + name = "sample_range" + + def convert( + self, value: Any, param: Optional[click.Parameter], ctx: Optional[click.Context] + ) -> Optional[tuple[int, int]]: + if isinstance(value, bool): + return DEFAULT_SAMPLE if value else None + if isinstance(value, (list, tuple)) and len(value) == 2: + value = f"{value[0]}-{value[1]}" + + match = re.match(r"^(\d+)-(\d+)$", str(value)) + if not match or int(match.group(1)) >= int(match.group(2)): + self.fail(f"'{value}' is not a valid sample window. Use format: START-END (e.g., 60-120)", param, ctx) + + return (int(match.group(1)), int(match.group(2))) + + SEASON_RANGE = SeasonRange() LANGUAGE_RANGE = LanguageRange() QUALITY_LIST = QualityList() @@ -512,3 +533,4 @@ AUDIO_CODEC_LIST = AudioCodecList(Audio.Codec) VIDEO_CODEC_LIST = MultipleVideoCodecChoice(Video.Codec) SUBTITLE_CODEC = SubtitleCodecChoice(Subtitle.Codec) SLOW_DELAY_RANGE = SlowDelayRange() +SAMPLE_RANGE = SampleRange()