perf(sample): download only the subtitle segments around the window

A segmented subtitle track took every segment in a --sample run. It now takes the window plus one segment on each side.

docs/guide/downloading.md

@@ -726,7 +726,8 @@ By default, `mkvmerge` muxes the tracks of a movie or an episode into a single M
     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.
+    A sidecar subtitle file with `sidecar_format: original` is not trimmed. It holds the segments around the
+    window, or the whole file when the track is one file.
 
 !!! note "Music is never muxed"
     A song skips the muxer. unshackle keeps the container that the service

tests/test_sample.py

@@ -54,10 +54,19 @@ def test_a_video_keeps_one_segment_after_the_window() -> None:
     assert make_audio().sample_slice(STARTS, 1000, (0, 60)) == slice(0, 15)
 
 
-def test_a_subtitle_keeps_every_segment() -> None:
+def test_a_subtitle_keeps_one_segment_on_each_side_of_the_window() -> 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
+    video = make_video()
+    video.sample_tail = 0
+    keep = video.sample_slice(STARTS, 1000, (62, 120))
+    assert subtitle.sample_slice(STARTS, 1000, (62, 120)) == slice(keep.start - 1, keep.stop + 1)
+    assert subtitle.sample_start == 0, "cue times stay absolute, so the mux does not shift the track"
+    assert subtitle.sample_slice(STARTS, 1000, None) == slice(None)
+
+
+def test_a_subtitle_that_ends_before_the_window_keeps_its_last_segment() -> None:
+    subtitle = Subtitle(id_="s", url="https://x/s", codec=Subtitle.Codec.WebVTT, language="en")
+    assert subtitle.sample_slice([0, 4000], 1000, (120, 180)) == slice(0, None)
 
 
 def test_a_window_after_the_end_stops_the_run() -> None:
@@ -172,19 +181,19 @@ def served(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> str:
     server.shutdown()
 
 
-def test_an_hls_subtitle_downloads_whole_in_a_sample_run(
+def test_an_hls_subtitle_downloads_the_segments_around_the_window(
     tmp_path: Path, served: str, monkeypatch: pytest.MonkeyPatch
 ) -> None:
-    for n in range(4):
+    for n in range(8):
         (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))
+        + "".join(f"#EXTINF:10.0,\ns{n}.vtt\n" for n in range(8))
         + "#EXT-X-ENDLIST\n"
     )
-    monkeypatch.setattr(config, "sample", (12, 18))
+    monkeypatch.setattr(config, "sample", (32, 38))
     subtitle = Subtitle(
         id_="s", url=f"{served}/s.m3u8", codec=Subtitle.Codec.WebVTT, language="en", descriptor=Subtitle.Descriptor.HLS
     )
@@ -192,7 +201,8 @@ def test_an_hls_subtitle_downloads_whole_in_a_sample_run(
     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))
+    assert [n for n in range(8) if f"cue {n}" in text] == [2, 3, 4]
+    assert "00:00:31.000 --> 00:00:32.000" in text, "cue times stay absolute"
 
 
 def test_subtitle_trim_keeps_the_cues_of_the_window_from_zero(tmp_path: Path) -> None:

unshackle/core/tracks/subtitle.py

@@ -11,6 +11,7 @@ from io import BytesIO
 from pathlib import Path
 from typing import Any, Callable, Iterable, Optional, Sequence, Union
 
+import click
 import pycaption
 import requests
 from construct import Container
@@ -1004,8 +1005,19 @@ class Subtitle(Track):
     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)
+        """Return the segments of ``window`` plus one segment on each side.
+
+        Cue times in a segment are absolute, so ``sample_start`` stays zero and ``trim`` makes the cut.
+        A track that ends before the window keeps only its last segment.
+        """
+        try:
+            keep = super().sample_slice(starts, timescale, window)
+        except click.ClickException:
+            keep = slice(len(starts) - 1, None)
+        self.sample_start = 0.0
+        if keep.start is None:
+            return keep
+        return slice(max(keep.start - 1, 0), None if keep.stop is None else keep.stop + 1)
 
     def trim(self, start: float, end: float) -> bool:
         """Keep the cues between ``start`` and ``end`` seconds and move them to start at zero.