From d7476ac6c25f87be95ae3e07207f0cd29defe111 Mon Sep 17 00:00:00 2001 From: imSp4rky Date: Sat, 14 Nov 2026 22:15:49 -0700 Subject: [PATCH] refactor(tracks): simplify the dynamic HDR probe and read less for TS audio --- tests/core/test_hdr10plus_probe.py | 37 ++++++++++++---------- unshackle/core/tracks/audio.py | 2 +- unshackle/core/utils/hdr10plus.py | 51 +++++++++++++++++------------- 3 files changed, 50 insertions(+), 40 deletions(-) diff --git a/tests/core/test_hdr10plus_probe.py b/tests/core/test_hdr10plus_probe.py index 31919c28..9936bbb6 100644 --- a/tests/core/test_hdr10plus_probe.py +++ b/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_hdr10plus(fragment(nal(32, b"vps"))) is None - assert fragment_has_hdr10plus(b"\x47" * 188) is None + assert fragment_has_t35(box(b"moof", b"") + box(b"mdat", b"\xff" * 64), HDR10PLUS_T35) is None + assert fragment_has_t35(fragment(nal(32, b"vps")), HDR10PLUS_T35) 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( - hdr10plus, "fragment_has_hdr10plus", lambda data, t35=HDR10PLUS_T35: t35 == HDR10PLUS_T35 and data == b"+" - ) + monkeypatch.setattr(hdr10plus, "fragment_has_t35", lambda data, 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: - fetched.append(track.id) + def fake_get(url: str, session: Any, byte_range: str) -> bytes: + 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 diff --git a/unshackle/core/tracks/audio.py b/unshackle/core/tracks/audio.py index 18268fcc..de188ba9 100644 --- a/unshackle/core/tracks/audio.py +++ b/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 diff --git a/unshackle/core/utils/hdr10plus.py b/unshackle/core/utils/hdr10plus.py index 1aa16926..30188b8b 100644 --- a/unshackle/core/utils/hdr10plus.py +++ b/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: - """Whether an HEVC SEI NAL holds an ST 2094-40 message, or the T.35 message that ``t35`` starts.""" +def sei_has_t35(nal: bytes, t35: bytes) -> bool: + """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]: - """Whether the first sample of an fMP4 fragment carries HDR10+ SEI, or None when the bytes cannot tell. - - ``t35`` picks another T.35 message, such as ``CUVA_T35`` for HDR Vivid. +def fragment_has_t35(data: bytes, t35: bytes) -> Optional[bool]: + """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. 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]: - """The range the first segment proves (HDR10P or VIVID), False for plain HDR10, None when it cannot tell.""" - found: Union[Video.Range, bool, None] = None + def proven(data: Optional[bytes]) -> Optional["Video.Range"]: + if not data: + 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) - 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 + found = proven(first_fragment(video, video.session or session)) 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})")