feat(dl): add --sample and --include-sample

--sample downloads only a time window of a title (default 120-180 s) and cuts the mux to it. --include-sample cuts a 60 s sample from a finished file. Both write to a Sample folder with a SAMPLE tag on the filename
This commit is contained in:
imSp4rky
2026-10-11 14:53:50 -06:00
parent 0a7c76561c
commit ecdfb667c9
20 changed files with 958 additions and 30 deletions
+2
View File
@@ -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://<resolver>@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. |
+2
View File
@@ -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). |
+14
View File
@@ -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.
+16
View File
@@ -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.
+456
View File
@@ -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 = """
<MPD xmlns="urn:mpeg:dash:schema:mpd:2011" type="static" mediaPresentationDuration="PT400S">
<Period id="0" start="PT0S">
<AdaptationSet mimeType="audio/mp4">
<SegmentTemplate timescale="1" media="a/$Number$-$Time$.mp4" startNumber="5">
<SegmentTimeline><S t="1000" d="4" r="99"/></SegmentTimeline>
</SegmentTemplate>
<Representation id="a" bandwidth="1"/>
</AdaptationSet>
</Period>
</MPD>
"""
DURATION_MPD = """
<MPD xmlns="urn:mpeg:dash:schema:mpd:2011" type="static" mediaPresentationDuration="PT400S">
<Period id="0" start="PT0S" duration="PT400S">
<AdaptationSet mimeType="audio/mp4">
<SegmentTemplate timescale="25" duration="100" media="a/$Number$.mp4" startNumber="5"/>
<Representation id="a" bandwidth="1"/>
</AdaptationSet>
</Period>
</MPD>
"""
LIST_MPD = """
<MPD xmlns="urn:mpeg:dash:schema:mpd:2011" type="static" mediaPresentationDuration="PT400S">
<Period id="0" start="PT0S">
<AdaptationSet mimeType="audio/mp4">
<Representation id="a" bandwidth="1">
<BaseURL>https://example.com/a/</BaseURL>
<SegmentList timescale="1" duration="4">{urls}</SegmentList>
</Representation>
</AdaptationSet>
</Period>
</MPD>
""".format(urls="".join(f'<SegmentURL media="{n}.mp4"/>' for n in range(100)))
BASE_MPD = """
<MPD xmlns="urn:mpeg:dash:schema:mpd:2011" type="static" mediaPresentationDuration="PT400S">
<Period id="0" start="PT0S">
<AdaptationSet mimeType="audio/mp4">
<Representation id="a" bandwidth="1"><BaseURL>https://example.com/a.mp4</BaseURL></Representation>
</AdaptationSet>
</Period>
</MPD>
"""
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 = """
<MPD xmlns="urn:mpeg:dash:schema:mpd:2011" type="static" mediaPresentationDuration="PT400S">
<Period id="0" start="PT0S">
<AdaptationSet mimeType="audio/mp4">
<Representation id="a" bandwidth="1">
<BaseURL>https://example.com/a.mp4</BaseURL>
<SegmentBase indexRange="800-{last}"><Initialization range="0-799"/></SegmentBase>
</Representation>
</AdaptationSet>
</Period>
</MPD>
"""
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
+131 -5
View File
@@ -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.<TAG>.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)
+11
View File
@@ -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),
+12
View File
@@ -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()
+8
View File
@@ -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)
+1
View File
@@ -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 {}
+2
View File
@@ -3,6 +3,8 @@ from typing import TypeVar, Union
DOWNLOAD_CANCELLED = Event()
DOWNLOAD_LICENCE_ONLY = Event()
DEFAULT_SAMPLE = (120, 180)
DOWNLOAD_ALL_DRM = Event()
+65 -10
View File
@@ -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))
+13 -6
View File
@@ -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),
+8
View File
@@ -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:
+6 -4
View File
@@ -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])
+27 -1
View File
@@ -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.
+122 -2
View File
@@ -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)
+38 -2
View File
@@ -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
+2
View File
@@ -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"
+22
View File
@@ -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()