From 3ba7740dd841ebcfe8f47612eac30d3b470fa93d Mon Sep 17 00:00:00 2001 From: pukkandan Date: Wed, 21 Jul 2021 22:58:43 +0530 Subject: [PATCH] [downloader] Pass `info_dict` to `progress_hook`s --- yt_dlp/YoutubeDL.py | 1 + yt_dlp/downloader/common.py | 12 +++++++++--- yt_dlp/downloader/dash.py | 2 +- yt_dlp/downloader/external.py | 2 +- yt_dlp/downloader/f4m.py | 4 ++-- yt_dlp/downloader/fragment.py | 15 ++++++++------- yt_dlp/downloader/hls.py | 2 +- yt_dlp/downloader/http.py | 6 +++--- yt_dlp/downloader/ism.py | 4 ++-- yt_dlp/downloader/mhtml.py | 4 ++-- yt_dlp/downloader/rtmp.py | 6 +++--- yt_dlp/downloader/rtsp.py | 2 +- yt_dlp/downloader/youtube_live_chat.py | 4 ++-- 13 files changed, 36 insertions(+), 28 deletions(-) diff --git a/yt_dlp/YoutubeDL.py b/yt_dlp/YoutubeDL.py index 2a7c05374e..0cba95bb66 100644 --- a/yt_dlp/YoutubeDL.py +++ b/yt_dlp/YoutubeDL.py @@ -322,6 +322,7 @@ class YoutubeDL(object): progress, with a dictionary with the entries * status: One of "downloading", "error", or "finished". Check this first and ignore unknown values. + * info_dict: The extracted info_dict If status is one of "downloading", or "finished", the following properties may also be present: diff --git a/yt_dlp/downloader/common.py b/yt_dlp/downloader/common.py index 9bf7eef3bd..9f0d3c7bf3 100644 --- a/yt_dlp/downloader/common.py +++ b/yt_dlp/downloader/common.py @@ -1,5 +1,6 @@ from __future__ import division, unicode_literals +import copy import os import re import sys @@ -360,7 +361,7 @@ def download(self, filename, info_dict, subtitle=False): 'filename': filename, 'status': 'finished', 'total_bytes': os.path.getsize(encodeFilename(filename)), - }) + }, info_dict) return True, False if subtitle is False: @@ -388,9 +389,14 @@ def real_download(self, filename, info_dict): """Real download process. Redefine in subclasses.""" raise NotImplementedError('This method must be implemented by subclasses') - def _hook_progress(self, status): + def _hook_progress(self, status, info_dict): + if not self._progress_hooks: + return + info_dict = dict(info_dict) + for key in ('__original_infodict', '__postprocessors'): + info_dict.pop(key, None) for ph in self._progress_hooks: - ph(status) + ph({**status, 'info_dict': copy.deepcopy(info_dict)}) def add_progress_hook(self, ph): # See YoutubeDl.py (search for progress_hooks) for a description of diff --git a/yt_dlp/downloader/dash.py b/yt_dlp/downloader/dash.py index aa7728efd9..9dae6b9bd0 100644 --- a/yt_dlp/downloader/dash.py +++ b/yt_dlp/downloader/dash.py @@ -29,7 +29,7 @@ def real_download(self, filename, info_dict): if real_downloader: self._prepare_external_frag_download(ctx) else: - self._prepare_and_start_frag_download(ctx) + self._prepare_and_start_frag_download(ctx, info_dict) fragments_to_download = [] frag_index = 0 diff --git a/yt_dlp/downloader/external.py b/yt_dlp/downloader/external.py index bfe444e882..d0ee745b32 100644 --- a/yt_dlp/downloader/external.py +++ b/yt_dlp/downloader/external.py @@ -67,7 +67,7 @@ def real_download(self, filename, info_dict): 'downloaded_bytes': fsize, 'total_bytes': fsize, }) - self._hook_progress(status) + self._hook_progress(status, info_dict) return True else: self.to_stderr('\n') diff --git a/yt_dlp/downloader/f4m.py b/yt_dlp/downloader/f4m.py index 3eb406152c..9da2776d92 100644 --- a/yt_dlp/downloader/f4m.py +++ b/yt_dlp/downloader/f4m.py @@ -380,7 +380,7 @@ def real_download(self, filename, info_dict): base_url_parsed = compat_urllib_parse_urlparse(base_url) - self._start_frag_download(ctx) + self._start_frag_download(ctx, info_dict) frag_index = 0 while fragments_list: @@ -434,6 +434,6 @@ def real_download(self, filename, info_dict): msg = 'Missed %d fragments' % (fragments_list[0][1] - (frag_i + 1)) self.report_warning(msg) - self._finish_frag_download(ctx) + self._finish_frag_download(ctx, info_dict) return True diff --git a/yt_dlp/downloader/fragment.py b/yt_dlp/downloader/fragment.py index 8e211c766b..88238b64d4 100644 --- a/yt_dlp/downloader/fragment.py +++ b/yt_dlp/downloader/fragment.py @@ -83,9 +83,9 @@ def _prepare_url(self, info_dict, url): headers = info_dict.get('http_headers') return sanitized_Request(url, None, headers) if headers else url - def _prepare_and_start_frag_download(self, ctx): + def _prepare_and_start_frag_download(self, ctx, info_dict): self._prepare_frag_download(ctx) - self._start_frag_download(ctx) + self._start_frag_download(ctx, info_dict) def __do_ytdl_file(self, ctx): return not ctx['live'] and not ctx['tmpfilename'] == '-' and not self.params.get('_no_ytdl_file') @@ -219,7 +219,7 @@ def _prepare_frag_download(self, ctx): 'complete_frags_downloaded_bytes': resume_len, }) - def _start_frag_download(self, ctx): + def _start_frag_download(self, ctx, info_dict): resume_len = ctx['complete_frags_downloaded_bytes'] total_frags = ctx['total_frags'] # This dict stores the download progress, it's updated by the progress @@ -248,6 +248,7 @@ def frag_progress_hook(s): time_now = time.time() state['elapsed'] = time_now - start frag_total_bytes = s.get('total_bytes') or 0 + s['fragment_info_dict'] = s.pop('info_dict', {}) if not ctx['live']: estimated_size = ( (ctx['complete_frags_downloaded_bytes'] + frag_total_bytes) @@ -270,13 +271,13 @@ def frag_progress_hook(s): state['speed'] = s.get('speed') or ctx.get('speed') ctx['speed'] = state['speed'] ctx['prev_frag_downloaded_bytes'] = frag_downloaded_bytes - self._hook_progress(state) + self._hook_progress(state, info_dict) ctx['dl'].add_progress_hook(frag_progress_hook) return start - def _finish_frag_download(self, ctx): + def _finish_frag_download(self, ctx, info_dict): ctx['dest_stream'].close() if self.__do_ytdl_file(ctx): ytdl_filename = encodeFilename(self.ytdl_filename(ctx['filename'])) @@ -303,7 +304,7 @@ def _finish_frag_download(self, ctx): 'filename': ctx['filename'], 'status': 'finished', 'elapsed': elapsed, - }) + }, info_dict) def _prepare_external_frag_download(self, ctx): if 'live' not in ctx: @@ -421,5 +422,5 @@ def _download_fragment(fragment): if not result: return False - self._finish_frag_download(ctx) + self._finish_frag_download(ctx, info_dict) return True diff --git a/yt_dlp/downloader/hls.py b/yt_dlp/downloader/hls.py index 52433e5afe..64637badf6 100644 --- a/yt_dlp/downloader/hls.py +++ b/yt_dlp/downloader/hls.py @@ -133,7 +133,7 @@ def is_ad_fragment_end(s): if real_downloader: self._prepare_external_frag_download(ctx) else: - self._prepare_and_start_frag_download(ctx) + self._prepare_and_start_frag_download(ctx, info_dict) extra_state = ctx.setdefault('extra_state', {}) diff --git a/yt_dlp/downloader/http.py b/yt_dlp/downloader/http.py index 15eb54aab0..9830f9e273 100644 --- a/yt_dlp/downloader/http.py +++ b/yt_dlp/downloader/http.py @@ -177,7 +177,7 @@ def establish_connection(): 'status': 'finished', 'downloaded_bytes': ctx.resume_len, 'total_bytes': ctx.resume_len, - }) + }, info_dict) raise SucceedDownload() else: # The length does not match, we start the download over @@ -310,7 +310,7 @@ def retry(e): 'eta': eta, 'speed': speed, 'elapsed': now - ctx.start_time, - }) + }, info_dict) if data_len is not None and byte_counter == data_len: break @@ -357,7 +357,7 @@ def retry(e): 'filename': ctx.filename, 'status': 'finished', 'elapsed': time.time() - ctx.start_time, - }) + }, info_dict) return True diff --git a/yt_dlp/downloader/ism.py b/yt_dlp/downloader/ism.py index 07d74aef0b..09516abe56 100644 --- a/yt_dlp/downloader/ism.py +++ b/yt_dlp/downloader/ism.py @@ -246,7 +246,7 @@ def real_download(self, filename, info_dict): 'total_frags': len(segments), } - self._prepare_and_start_frag_download(ctx) + self._prepare_and_start_frag_download(ctx, info_dict) extra_state = ctx.setdefault('extra_state', { 'ism_track_written': False, @@ -284,6 +284,6 @@ def real_download(self, filename, info_dict): self.report_error('giving up after %s fragment retries' % fragment_retries) return False - self._finish_frag_download(ctx) + self._finish_frag_download(ctx, info_dict) return True diff --git a/yt_dlp/downloader/mhtml.py b/yt_dlp/downloader/mhtml.py index 81d95c7cbe..b75db18a8a 100644 --- a/yt_dlp/downloader/mhtml.py +++ b/yt_dlp/downloader/mhtml.py @@ -122,7 +122,7 @@ def real_download(self, filename, info_dict): 'total_frags': len(fragments), } - self._prepare_and_start_frag_download(ctx) + self._prepare_and_start_frag_download(ctx, info_dict) extra_state = ctx.setdefault('extra_state', { 'header_written': False, @@ -198,5 +198,5 @@ def real_download(self, filename, info_dict): ctx['dest_stream'].write( b'--%b--\r\n\r\n' % frag_boundary.encode('us-ascii')) - self._finish_frag_download(ctx) + self._finish_frag_download(ctx, info_dict) return True diff --git a/yt_dlp/downloader/rtmp.py b/yt_dlp/downloader/rtmp.py index 99158e6215..6dca64725d 100644 --- a/yt_dlp/downloader/rtmp.py +++ b/yt_dlp/downloader/rtmp.py @@ -66,7 +66,7 @@ def run_rtmpdump(args): 'eta': eta, 'elapsed': time_now - start, 'speed': speed, - }) + }, info_dict) cursor_in_new_line = False else: # no percent for live streams @@ -82,7 +82,7 @@ def run_rtmpdump(args): 'status': 'downloading', 'elapsed': time_now - start, 'speed': speed, - }) + }, info_dict) cursor_in_new_line = False elif self.params.get('verbose', False): if not cursor_in_new_line: @@ -208,7 +208,7 @@ def run_rtmpdump(args): 'filename': filename, 'status': 'finished', 'elapsed': time.time() - started, - }) + }, info_dict) return True else: self.to_stderr('\n') diff --git a/yt_dlp/downloader/rtsp.py b/yt_dlp/downloader/rtsp.py index 4ce2fafff3..7815d59d97 100644 --- a/yt_dlp/downloader/rtsp.py +++ b/yt_dlp/downloader/rtsp.py @@ -39,7 +39,7 @@ def real_download(self, filename, info_dict): 'total_bytes': fsize, 'filename': filename, 'status': 'finished', - }) + }, info_dict) return True else: self.to_stderr('\n') diff --git a/yt_dlp/downloader/youtube_live_chat.py b/yt_dlp/downloader/youtube_live_chat.py index 5e05426e63..2dc6ff954c 100644 --- a/yt_dlp/downloader/youtube_live_chat.py +++ b/yt_dlp/downloader/youtube_live_chat.py @@ -140,7 +140,7 @@ def download_and_parse_fragment(url, frag_index, request_data=None, headers=None self.report_error('giving up after %s fragment retries' % fragment_retries) return False, None, None, None - self._prepare_and_start_frag_download(ctx) + self._prepare_and_start_frag_download(ctx, info_dict) success, raw_fragment = dl_fragment(info_dict['url']) if not success: @@ -196,7 +196,7 @@ def download_and_parse_fragment(url, frag_index, request_data=None, headers=None if test: break - self._finish_frag_download(ctx) + self._finish_frag_download(ctx, info_dict) return True @staticmethod