From bd4d1ea398ab6118b5d16ff3e19562d477ff04be Mon Sep 17 00:00:00 2001 From: pukkandan Date: Thu, 24 Jun 2021 22:23:33 +0530 Subject: [PATCH] [cleanup] Minor refactoring of `fragment` --- yt_dlp/downloader/dash.py | 4 +- yt_dlp/downloader/fragment.py | 10 +-- yt_dlp/downloader/hls.py | 140 +++++++++++++++++----------------- 3 files changed, 76 insertions(+), 78 deletions(-) diff --git a/yt_dlp/downloader/dash.py b/yt_dlp/downloader/dash.py index ce4001a09..aa7728efd 100644 --- a/yt_dlp/downloader/dash.py +++ b/yt_dlp/downloader/dash.py @@ -58,5 +58,5 @@ class DashSegmentsFD(FragmentFD): # for ph in self._progress_hooks: # fd.add_progress_hook(ph) return fd.real_download(filename, info_copy) - else: - return self.download_and_append_fragments(ctx, fragments_to_download, info_dict) + + return self.download_and_append_fragments(ctx, fragments_to_download, info_dict) diff --git a/yt_dlp/downloader/fragment.py b/yt_dlp/downloader/fragment.py index c8d4d33aa..c499e5e2b 100644 --- a/yt_dlp/downloader/fragment.py +++ b/yt_dlp/downloader/fragment.py @@ -328,8 +328,7 @@ class FragmentFD(FileDownloader): def download_and_append_fragments(self, ctx, fragments, info_dict, pack_func=None): fragment_retries = self.params.get('fragment_retries', 0) - skip_unavailable_fragments = self.params.get('skip_unavailable_fragments', True) - test = self.params.get('test', False) + is_fatal = (lambda idx: idx == 0) if self.params.get('skip_unavailable_fragments', True) else (lambda _: True) if not pack_func: pack_func = lambda frag_content, _: frag_content @@ -341,7 +340,7 @@ class FragmentFD(FileDownloader): headers['Range'] = 'bytes=%d-%d' % (byte_range['start'], byte_range['end'] - 1) # Never skip the first fragment - fatal = (fragment.get('index') or frag_index) == 0 or not skip_unavailable_fragments + fatal = is_fatal(fragment.get('index') or (frag_index - 1)) count, frag_content = 0, None while count <= fragment_retries: try: @@ -382,14 +381,13 @@ class FragmentFD(FileDownloader): # Don't decrypt the content in tests since the data is explicitly truncated and it's not to a valid block # size (see https://github.com/ytdl-org/youtube-dl/pull/27660). Tests only care that the correct data downloaded, # not what it decrypts to. - if test: + if self.params.get('test', False): return frag_content return AES.new(decrypt_info['KEY'], AES.MODE_CBC, iv).decrypt(frag_content) def append_fragment(frag_content, frag_index, ctx): if not frag_content: - fatal = frag_index == 1 or not skip_unavailable_fragments - if not fatal: + if not is_fatal(frag_index - 1): self.report_skip_fragment(frag_index) return True else: diff --git a/yt_dlp/downloader/hls.py b/yt_dlp/downloader/hls.py index fdc4f7aa8..52433e5af 100644 --- a/yt_dlp/downloader/hls.py +++ b/yt_dlp/downloader/hls.py @@ -251,74 +251,74 @@ class HlsFD(FragmentFD): # for ph in self._progress_hooks: # fd.add_progress_hook(ph) return fd.real_download(filename, info_copy) + + if is_webvtt: + def pack_fragment(frag_content, frag_index): + output = io.StringIO() + adjust = 0 + for block in webvtt.parse_fragment(frag_content): + if isinstance(block, webvtt.CueBlock): + block.start += adjust + block.end += adjust + + dedup_window = extra_state.setdefault('webvtt_dedup_window', []) + cue = block.as_json + + # skip the cue if an identical one appears + # in the window of potential duplicates + # and prune the window of unviable candidates + i = 0 + skip = True + while i < len(dedup_window): + window_cue = dedup_window[i] + if window_cue == cue: + break + if window_cue['end'] >= cue['start']: + i += 1 + continue + del dedup_window[i] + else: + skip = False + + if skip: + continue + + # add the cue to the window + dedup_window.append(cue) + elif isinstance(block, webvtt.Magic): + # take care of MPEG PES timestamp overflow + if block.mpegts is None: + block.mpegts = 0 + extra_state.setdefault('webvtt_mpegts_adjust', 0) + block.mpegts += extra_state['webvtt_mpegts_adjust'] << 33 + if block.mpegts < extra_state.get('webvtt_mpegts_last', 0): + extra_state['webvtt_mpegts_adjust'] += 1 + block.mpegts += 1 << 33 + extra_state['webvtt_mpegts_last'] = block.mpegts + + if frag_index == 1: + extra_state['webvtt_mpegts'] = block.mpegts or 0 + extra_state['webvtt_local'] = block.local or 0 + # XXX: block.local = block.mpegts = None ? + else: + if block.mpegts is not None and block.local is not None: + adjust = ( + (block.mpegts - extra_state.get('webvtt_mpegts', 0)) + - (block.local - extra_state.get('webvtt_local', 0)) + ) + continue + elif isinstance(block, webvtt.HeaderBlock): + if frag_index != 1: + # XXX: this should probably be silent as well + # or verify that all segments contain the same data + self.report_warning(bug_reports_message( + 'Discarding a %s block found in the middle of the stream; ' + 'if the subtitles display incorrectly,' + % (type(block).__name__))) + continue + block.write_into(output) + + return output.getvalue().encode('utf-8') else: - if is_webvtt: - def pack_fragment(frag_content, frag_index): - output = io.StringIO() - adjust = 0 - for block in webvtt.parse_fragment(frag_content): - if isinstance(block, webvtt.CueBlock): - block.start += adjust - block.end += adjust - - dedup_window = extra_state.setdefault('webvtt_dedup_window', []) - cue = block.as_json - - # skip the cue if an identical one appears - # in the window of potential duplicates - # and prune the window of unviable candidates - i = 0 - skip = True - while i < len(dedup_window): - window_cue = dedup_window[i] - if window_cue == cue: - break - if window_cue['end'] >= cue['start']: - i += 1 - continue - del dedup_window[i] - else: - skip = False - - if skip: - continue - - # add the cue to the window - dedup_window.append(cue) - elif isinstance(block, webvtt.Magic): - # take care of MPEG PES timestamp overflow - if block.mpegts is None: - block.mpegts = 0 - extra_state.setdefault('webvtt_mpegts_adjust', 0) - block.mpegts += extra_state['webvtt_mpegts_adjust'] << 33 - if block.mpegts < extra_state.get('webvtt_mpegts_last', 0): - extra_state['webvtt_mpegts_adjust'] += 1 - block.mpegts += 1 << 33 - extra_state['webvtt_mpegts_last'] = block.mpegts - - if frag_index == 1: - extra_state['webvtt_mpegts'] = block.mpegts or 0 - extra_state['webvtt_local'] = block.local or 0 - # XXX: block.local = block.mpegts = None ? - else: - if block.mpegts is not None and block.local is not None: - adjust = ( - (block.mpegts - extra_state.get('webvtt_mpegts', 0)) - - (block.local - extra_state.get('webvtt_local', 0)) - ) - continue - elif isinstance(block, webvtt.HeaderBlock): - if frag_index != 1: - # XXX: this should probably be silent as well - # or verify that all segments contain the same data - self.report_warning(bug_reports_message( - 'Discarding a %s block found in the middle of the stream; ' - 'if the subtitles display incorrectly,' - % (type(block).__name__))) - continue - block.write_into(output) - - return output.getvalue().encode('utf-8') - else: - pack_fragment = None - return self.download_and_append_fragments(ctx, fragments, info_dict, pack_fragment) + pack_fragment = None + return self.download_and_append_fragments(ctx, fragments, info_dict, pack_fragment)