1
0
mirror of https://github.com/yt-dlp/yt-dlp.git synced 2025-06-28 01:18:30 +00:00
yt-dlp/yt_dlp/extractor/youtube/_streaming/ump.py
2025-06-21 11:15:25 +12:00

117 lines
3.1 KiB
Python

import dataclasses
import enum
import io
class UMPPartId(enum.IntEnum):
UNKNOWN = -1
ONESIE_HEADER = 10
ONESIE_DATA = 11
ONESIE_ENCRYPTED_MEDIA = 12
MEDIA_HEADER = 20
MEDIA = 21
MEDIA_END = 22
LIVE_METADATA = 31
HOSTNAME_CHANGE_HINT = 32
LIVE_METADATA_PROMISE = 33
LIVE_METADATA_PROMISE_CANCELLATION = 34
NEXT_REQUEST_POLICY = 35
USTREAMER_VIDEO_AND_FORMAT_DATA = 36
FORMAT_SELECTION_CONFIG = 37
USTREAMER_SELECTED_MEDIA_STREAM = 38
FORMAT_INITIALIZATION_METADATA = 42
SABR_REDIRECT = 43
SABR_ERROR = 44
SABR_SEEK = 45
RELOAD_PLAYER_RESPONSE = 46
PLAYBACK_START_POLICY = 47
ALLOWED_CACHED_FORMATS = 48
START_BW_SAMPLING_HINT = 49
PAUSE_BW_SAMPLING_HINT = 50
SELECTABLE_FORMATS = 51
REQUEST_IDENTIFIER = 52
REQUEST_CANCELLATION_POLICY = 53
ONESIE_PREFETCH_REJECTION = 54
TIMELINE_CONTEXT = 55
REQUEST_PIPELINING = 56
SABR_CONTEXT_UPDATE = 57
STREAM_PROTECTION_STATUS = 58
SABR_CONTEXT_SENDING_POLICY = 59
LAWNMOWER_POLICY = 60
SABR_ACK = 61
END_OF_TRACK = 62
CACHE_LOAD_POLICY = 63
LAWNMOWER_MESSAGING_POLICY = 64
PREWARM_CONNECTION = 65
PLAYBACK_DEBUG_INFO = 66
SNACKBAR_MESSAGE = 67
@classmethod
def _missing_(cls, value):
return cls.UNKNOWN
@dataclasses.dataclass
class UMPPart:
part_id: UMPPartId
size: int
data: io.BufferedIOBase
class UMPDecoder:
def __init__(self, fp: io.BufferedIOBase):
self.fp = fp
def iter_parts(self):
while not self.fp.closed:
part_type = read_varint(self.fp)
if part_type == -1 and not self.fp.closed:
self.fp.close()
if self.fp.closed:
break
part_size = read_varint(self.fp)
if part_size == -1 and not self.fp.closed:
self.fp.close()
if self.fp.closed:
raise EOFError('Unexpected EOF while reading part size')
part_data = self.fp.read(part_size)
# In the future, we could allow streaming the part data.
# But we will need to ensure that each part is completely read before continuing.
yield UMPPart(UMPPartId(part_type), part_size, io.BytesIO(part_data))
def read_varint(fp: io.BufferedIOBase) -> int:
# https://web.archive.org/web/20250430054327/https://github.com/gsuberland/UMP_Format/blob/main/UMP_Format.md
# https://web.archive.org/web/20250429151021/https://github.com/davidzeng0/innertube/blob/main/googlevideo/ump.md
byte = fp.read(1)
if not byte:
# Expected EOF
return -1
prefix = byte[0]
size = varint_size(prefix)
result = 0
shift = 0
if size != 5:
shift = 8 - size
mask = (1 << shift) - 1
result |= prefix & mask
for _ in range(1, size):
next_byte = fp.read(1)
if not next_byte:
return -1
byte_int = next_byte[0]
result |= byte_int << shift
shift += 8
return result
def varint_size(byte: int) -> int:
return 1 if byte < 128 else 2 if byte < 192 else 3 if byte < 224 else 4 if byte < 240 else 5