refactor(tracks): simplify the dynamic HDR probe and read less for TS audio

tests/core/test_hdr10plus_probe.py

@@ -12,7 +12,7 @@ from unshackle.core.manifests.snapshot import match_track
 from unshackle.core.remote_service import RemoteService
 from unshackle.core.tracks import Audio, Video
 from unshackle.core.utils import hdr10plus
-from unshackle.core.utils.hdr10plus import HDR10PLUS_T35, TS_PACKET, fragment_has_hdr10plus
+from unshackle.core.utils.hdr10plus import HDR10PLUS_T35, TS_PACKET, fragment_has_t35
 
 
 def box(kind: bytes, body: bytes) -> bytes:
@@ -34,18 +34,18 @@ def fragment(*nals: bytes) -> bytes:
 
 def test_hdr10plus_sei_before_the_slice_is_found() -> None:
     data = fragment(nal(32, b"vps"), sei(137, b"\x00" * 24), sei(4, HDR10PLUS_T35 + b"\x00" * 8), nal(1, b"\xff" * 64))
-    assert fragment_has_hdr10plus(data) is True
+    assert fragment_has_t35(data, HDR10PLUS_T35) is True
 
 
 def test_static_metadata_only_is_plain_hdr10() -> None:
     data = fragment(nal(32, b"vps"), sei(137, b"\x00" * 24), sei(4, b"\xb5\x00\x31GA94"), nal(1, b"\xff" * 64))
-    assert fragment_has_hdr10plus(data) is False
+    assert fragment_has_t35(data, HDR10PLUS_T35) is False
 
 
 def test_encrypted_or_truncated_bytes_prove_nothing() -> None:
-    assert fragment_has_hdr10plus(box(b"moof", b"") + box(b"mdat", b"\xff" * 64)) is None
+    assert fragment_has_t35(box(b"moof", b"") + box(b"mdat", b"\xff" * 64), HDR10PLUS_T35) is None
-    assert fragment_has_hdr10plus(fragment(nal(32, b"vps"))) is None
+    assert fragment_has_t35(fragment(nal(32, b"vps")), HDR10PLUS_T35) is None
-    assert fragment_has_hdr10plus(b"\x47" * 188) is None
+    assert fragment_has_t35(b"\x47" * 188, HDR10PLUS_T35) is None
 
 
 def hevc(track_id: str, height: int, bitrate: int = 1, range_: Video.Range = Video.Range.HDR10) -> Any:
@@ -79,9 +79,7 @@ def probe(monkeypatch: pytest.MonkeyPatch, videos: list[Any]) -> list[str]:
         return video.id[-1].encode()
 
     monkeypatch.setattr(hdr10plus, "first_fragment", fake_fragment)
-    monkeypatch.setattr(
+    monkeypatch.setattr(hdr10plus, "fragment_has_t35", lambda data, t35: t35 == HDR10PLUS_T35 and data == b"+")
-        hdr10plus, "fragment_has_hdr10plus", lambda data, t35=HDR10PLUS_T35: t35 == HDR10PLUS_T35 and data == b"+"
-    )
     hdr10plus.probe_dynamic_hdr(videos, None, log=logging.getLogger("test"), workers=4)  # type: ignore[arg-type]
     return probed
 
@@ -157,7 +155,7 @@ def test_first_fragment_skips_the_header_boxes_of_a_one_file_track(monkeypatch:
 
     monkeypatch.setattr(hdr10plus, "first_segment", lambda track, session: ("https://cdn/v.mp4", 0))
     monkeypatch.setattr(hdr10plus, "ranged_get", fake_get)
-    assert fragment_has_hdr10plus(hdr10plus.first_fragment(hevc("v", 2160), None) or b"") is True  # type: ignore[arg-type]
+    assert fragment_has_t35(hdr10plus.first_fragment(hevc("v", 2160), None) or b"", HDR10PLUS_T35) is True  # type: ignore[arg-type]
 
 
 def test_first_segment_of_an_hls_playlist_keeps_its_byte_offset() -> None:
@@ -216,7 +214,7 @@ def test_a_fragment_inside_the_first_window_needs_one_request(monkeypatch: pytes
 
     monkeypatch.setattr(hdr10plus, "first_segment", lambda track, session: ("https://cdn/v.mp4", 0))
     monkeypatch.setattr(hdr10plus, "ranged_get", fake_get)
-    assert fragment_has_hdr10plus(hdr10plus.first_fragment(hevc("v", 2160), None) or b"") is True  # type: ignore[arg-type]
+    assert fragment_has_t35(hdr10plus.first_fragment(hevc("v", 2160), None) or b"", HDR10PLUS_T35) is True  # type: ignore[arg-type]
     assert len(requests) == 1
 
 
@@ -297,22 +295,27 @@ def test_ts_streams_skips_program_zero_and_survives_bad_lengths() -> None:
 
 
 def test_audio_without_a_codec_takes_it_from_the_ts_stream_type(monkeypatch: pytest.MonkeyPatch) -> None:
-    fetched: list[str] = []
+    fetched: list[tuple[str, str]] = []
 
-    def fake_fragment(track: Any, session: Any) -> bytes:
+    def fake_get(url: str, session: Any, byte_range: str) -> bytes:
-        fetched.append(track.id)
+        fetched.append((url, byte_range))
         return ts((0xD5, b""))
 
-    monkeypatch.setattr(hdr10plus, "first_fragment", fake_fragment)
+    monkeypatch.setattr(hdr10plus, "first_segment", lambda track, session: (track.id, 376))
+    monkeypatch.setattr(hdr10plus, "ranged_get", fake_get)
     vivid = SimpleNamespace(id="a1", codec=None, session=None)
     named = SimpleNamespace(id="a2", codec=Audio.Codec.AAC, session=None)
     hdr10plus.probe_audio_codecs([vivid, named], None, log=logging.getLogger("test"), workers=4)  # type: ignore[arg-type,list-item]
-    assert (vivid.codec, named.codec, fetched) == (Audio.Codec.AV3A, Audio.Codec.AAC, ["a1"])
+    assert (vivid.codec, named.codec, fetched) == (
+        Audio.Codec.AV3A,
+        Audio.Codec.AAC,
+        [("a1", f"bytes=376-{376 + 16 * TS_PACKET - 1}")],
+    )
     assert "[Audio Vivid]" in str(Audio(url="https://x/a.m3u8", language="zh", codec=Audio.Codec.AV3A))
 
 
 def test_an_avsa_registration_names_audio_vivid(monkeypatch: pytest.MonkeyPatch) -> None:
-    monkeypatch.setattr(hdr10plus, "first_fragment", lambda track, session: ts((0x06, b"AVSA")))
+    serve_segments(monkeypatch, {"a": ts((0x06, b"AVSA"))})
     track = SimpleNamespace(id="a", codec=None, session=None)
     hdr10plus.probe_audio_codecs([track], None, log=logging.getLogger("test"), workers=1)  # type: ignore[arg-type,list-item]
     assert track.codec == Audio.Codec.AV3A

unshackle/core/tracks/audio.py

@@ -101,7 +101,7 @@ class Audio(Track):
             raise ValueError(f"The MIME '{mime}' is not a supported Audio Codec")
 
         @staticmethod
-        def from_ts(stream_type: int, registration: bytes = b"") -> Optional[Audio.Codec]:
+        def from_ts(stream_type: int, registration: bytes) -> Optional[Audio.Codec]:
             """The codec an MPEG-TS PMT entry names by its registration descriptor or ``stream_type``, else None."""
             if registration == b"AVSA":
                 return Audio.Codec.AV3A

unshackle/core/utils/hdr10plus.py

@@ -21,6 +21,7 @@ PROBE_BYTES = 64 * 1024
 HDR10PLUS_T35 = b"\xb5\x00\x3c\x00\x01\x04"
 CUVA_T35 = b"\x26\x00\x04\x00\x05"
 TS_PACKET = 188
+TS_PROBE_BYTES = 16 * TS_PACKET
 T = TypeVar("T")
 R = TypeVar("R")
 
@@ -98,8 +99,8 @@ def sei_number(rbsp: bytes, i: int) -> tuple[Optional[int], int]:
     return value + rbsp[i], i + 1
 
 
-def sei_has_t35(nal: bytes, t35: bytes = HDR10PLUS_T35) -> bool:
+def sei_has_t35(nal: bytes, t35: bytes) -> bool:
-    """Whether an HEVC SEI NAL holds an ST 2094-40 message, or the T.35 message that ``t35`` starts."""
+    """Whether an HEVC SEI NAL holds the T.35 message that ``t35`` starts."""
     rbsp = remove_emulation_prevention(nal[2:])
     i = 0
     while i < len(rbsp) and rbsp[i] != 0x80:
@@ -113,10 +114,9 @@ def sei_has_t35(nal: bytes, t35: bytes = HDR10PLUS_T35) -> bool:
     return False
 
 
-def fragment_has_hdr10plus(data: bytes, t35: bytes = HDR10PLUS_T35) -> Optional[bool]:
+def fragment_has_t35(data: bytes, t35: bytes) -> Optional[bool]:
-    """Whether the first sample of an fMP4 fragment carries HDR10+ SEI, or None when the bytes cannot tell.
+    """Whether the first sample of an fMP4 fragment carries the T.35 SEI that ``t35`` starts, or None when the bytes
-
+    cannot tell. ``HDR10PLUS_T35`` finds HDR10+, ``CUVA_T35`` HDR Vivid.
-    ``t35`` picks another T.35 message, such as ``CUVA_T35`` for HDR Vivid.
 
     The walk ends at the first VCL NAL, because a prefix SEI comes before it. CENC and cbcs keep
     non-VCL NALs and every NAL length in the clear, so this needs no content key. A sample that is
@@ -241,26 +241,26 @@ def probe_dynamic_hdr(
     if any(ladders) and on_start:
         on_start()
 
-    def probe(video: "Video") -> Union["Video.Range", bool, None]:
+    def proven(data: Optional[bytes]) -> Optional["Video.Range"]:
-        """The range the first segment proves (HDR10P or VIVID), False for plain HDR10, None when it cannot tell."""
+        if not data:
-        found: Union[Video.Range, bool, None] = None
+            return None
+        if is_ts(data):
+            return Video.Range.VIVID if any(reg == b"cuvv" for _, reg in ts_streams(data) or []) else None
+        if fragment_has_t35(data, CUVA_T35):
+            return Video.Range.VIVID
+        return Video.Range.HDR10P if fragment_has_t35(data, HDR10PLUS_T35) else None
+
+    def probe(video: "Video") -> Optional["Video.Range"]:
+        """The range the first segment proves (HDR10P or VIVID), else None."""
         try:
-            data = first_fragment(video, video.session or session)
+            found = proven(first_fragment(video, video.session or session))
-            if data and is_ts(data):
-                if any(registration == b"cuvv" for _, registration in ts_streams(data) or []):
-                    found = Video.Range.VIVID
-            elif data and fragment_has_hdr10plus(data, CUVA_T35):
-                found = Video.Range.VIVID
-            elif data:
-                hdr10plus = fragment_has_hdr10plus(data)
-                found = Video.Range.HDR10P if hdr10plus else hdr10plus
         except Exception as e:
             log.debug(f"{video.id}: dynamic HDR probe failed ({e})")
             return None
         log.debug(f"{video.id}: dynamic HDR probe {found}")
         return found
 
-    def probe_all(tracks: list["Video"]) -> dict[int, Union["Video.Range", bool, None]]:
+    def probe_all(tracks: list["Video"]) -> dict[int, Optional["Video.Range"]]:
         return dict(zip(map(id, tracks), map_probes(probe, tracks, workers)))
 
     plans: list[tuple[list["Video"], list["Video"], list["Video"]]] = []
@@ -282,19 +282,26 @@ def probe_dynamic_hdr(
             label = found.get(id(video), whole)
             if video.range == Video.Range.DV:
                 video.hdr10plus = video.hdr10plus or label == Video.Range.HDR10P
-            elif isinstance(label, Video.Range) and label in PROBED_RANGES.get(video.range, ()):
+            elif label in PROBED_RANGES.get(video.range, ()):
                 video.range = label
 
 
 def probe_audio_codecs(
     audio: list["Audio"], session: Union[Session, RnetSession], *, log: logging.Logger, workers: int
 ) -> None:
-    """Name the codec of each audio track that has none from the PMT of its MPEG-TS first segment."""
+    """Name the codec of each audio track that has none from the PMT of its MPEG-TS first segment.
+
+    Only the first ``TS_PROBE_BYTES`` are read, because an HLS segment opens with its PAT and PMT.
+    """
     from unshackle.core.tracks import Audio
 
     def probe(track: "Audio") -> None:
         try:
-            data = first_fragment(track, track.session or session)
+            track_session = track.session or session
+            segment = first_segment(track, track_session)
+            data = segment and ranged_get(
+                segment[0], track_session, f"bytes={segment[1]}-{segment[1] + TS_PROBE_BYTES - 1}"
+            )
             streams = (ts_streams(data) if data and is_ts(data) else None) or []
         except Exception as e:
             log.debug(f"{track.id}: audio codec probe failed ({e})")