From ab8e5e516f38c3eab8947614e2347a2473e5dbbc Mon Sep 17 00:00:00 2001 From: Alex Merkel Date: Tue, 29 Dec 2020 16:03:07 +0100 Subject: [PATCH] Add post_hooks option to YoutubeDL.py (https://github.com/ytdl-org/youtube-dl/pull/27573) Authored by: alexmerkel --- Makefile | 3 +- devscripts/run_tests.bat | 2 +- devscripts/run_tests.sh | 2 +- test/test_post_hooks.py | 68 ++++++++++++++++++++++++++++++++++++++++ youtube_dlc/YoutubeDL.py | 17 ++++++++++ 5 files changed, 89 insertions(+), 3 deletions(-) create mode 100644 test/test_post_hooks.py diff --git a/Makefile b/Makefile index 317569e05..368c83585 100644 --- a/Makefile +++ b/Makefile @@ -50,7 +50,8 @@ offlinetest: codetest --exclude test_subtitles.py \ --exclude test_write_annotations.py \ --exclude test_youtube_lists.py \ - --exclude test_youtube_signature.py + --exclude test_youtube_signature.py \ + --exclude test_post_hooks.py tar: youtube-dlc.tar.gz diff --git a/devscripts/run_tests.bat b/devscripts/run_tests.bat index 79359b5a7..531af4066 100644 --- a/devscripts/run_tests.bat +++ b/devscripts/run_tests.bat @@ -1,7 +1,7 @@ @echo off rem Keep this list in sync with the `offlinetest` target in Makefile -set DOWNLOAD_TESTS="age_restriction^|download^|iqiyi_sdk_interpreter^|socks^|subtitles^|write_annotations^|youtube_lists^|youtube_signature" +set DOWNLOAD_TESTS="age_restriction^|download^|iqiyi_sdk_interpreter^|socks^|subtitles^|write_annotations^|youtube_lists^|youtube_signature^|post_hooks" if "%YTDL_TEST_SET%" == "core" ( set test_set="-I test_("%DOWNLOAD_TESTS%")\.py" diff --git a/devscripts/run_tests.sh b/devscripts/run_tests.sh index dd37a80f5..2fa7d16e2 100755 --- a/devscripts/run_tests.sh +++ b/devscripts/run_tests.sh @@ -1,7 +1,7 @@ #!/bin/bash # Keep this list in sync with the `offlinetest` target in Makefile -DOWNLOAD_TESTS="age_restriction|download|iqiyi_sdk_interpreter|socks|subtitles|write_annotations|youtube_lists|youtube_signature" +DOWNLOAD_TESTS="age_restriction|download|iqiyi_sdk_interpreter|socks|subtitles|write_annotations|youtube_lists|youtube_signature|post_hooks" test_set="" multiprocess_args="" diff --git a/test/test_post_hooks.py b/test/test_post_hooks.py new file mode 100644 index 000000000..d8d2b36c3 --- /dev/null +++ b/test/test_post_hooks.py @@ -0,0 +1,68 @@ +#!/usr/bin/env python + +from __future__ import unicode_literals + +import os +import sys +import unittest +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from test.helper import get_params, try_rm +import youtube_dl.YoutubeDL +from youtube_dl.utils import DownloadError + + +class YoutubeDL(youtube_dl.YoutubeDL): + def __init__(self, *args, **kwargs): + super(YoutubeDL, self).__init__(*args, **kwargs) + self.to_stderr = self.to_screen + + +TEST_ID = 'gr51aVj-mLg' +EXPECTED_NAME = 'gr51aVj-mLg' + + +class TestPostHooks(unittest.TestCase): + def setUp(self): + self.stored_name_1 = None + self.stored_name_2 = None + self.params = get_params({ + 'skip_download': False, + 'writeinfojson': False, + 'quiet': True, + 'verbose': False, + 'cachedir': False, + }) + self.files = [] + + def test_post_hooks(self): + self.params['post_hooks'] = [self.hook_one, self.hook_two] + ydl = YoutubeDL(self.params) + ydl.download([TEST_ID]) + self.assertEqual(self.stored_name_1, EXPECTED_NAME, 'Not the expected name from hook 1') + self.assertEqual(self.stored_name_2, EXPECTED_NAME, 'Not the expected name from hook 2') + + def test_post_hook_exception(self): + self.params['post_hooks'] = [self.hook_three] + ydl = YoutubeDL(self.params) + self.assertRaises(DownloadError, ydl.download, [TEST_ID]) + + def hook_one(self, filename): + self.stored_name_1, _ = os.path.splitext(os.path.basename(filename)) + self.files.append(filename) + + def hook_two(self, filename): + self.stored_name_2, _ = os.path.splitext(os.path.basename(filename)) + self.files.append(filename) + + def hook_three(self, filename): + self.files.append(filename) + raise Exception('Test exception for \'%s\'' % filename) + + def tearDown(self): + for f in self.files: + try_rm(f) + + +if __name__ == '__main__': + unittest.main() diff --git a/youtube_dlc/YoutubeDL.py b/youtube_dlc/YoutubeDL.py index fbd40cf73..3bae07764 100644 --- a/youtube_dlc/YoutubeDL.py +++ b/youtube_dlc/YoutubeDL.py @@ -252,6 +252,9 @@ class YoutubeDL(object): youtube_dlc/postprocessor/__init__.py for a list. as well as any further keyword arguments for the postprocessor. + post_hooks: A list of functions that get called as the final step + for each video file, after all postprocessors have been + called. The filename will be passed as the only argument. progress_hooks: A list of functions that get called on download progress, with a dictionary with the entries * status: One of "downloading", "error", or "finished". @@ -369,6 +372,7 @@ class YoutubeDL(object): self._ies = [] self._ies_instances = {} self._pps = [] + self._post_hooks = [] self._progress_hooks = [] self._download_retcode = 0 self._num_downloads = 0 @@ -472,6 +476,9 @@ class YoutubeDL(object): pp = pp_class(self, **compat_kwargs(pp_def)) self.add_post_processor(pp) + for ph in self.params.get('post_hooks', []): + self.add_post_hook(ph) + for ph in self.params.get('progress_hooks', []): self.add_progress_hook(ph) @@ -524,6 +531,10 @@ class YoutubeDL(object): self._pps.append(pp) pp.set_downloader(self) + def add_post_hook(self, ph): + """Add the post hook""" + self._post_hooks.append(ph) + def add_progress_hook(self, ph): """Add the progress hook (currently only for the file downloader)""" self._progress_hooks.append(ph) @@ -2199,6 +2210,12 @@ class YoutubeDL(object): except (PostProcessingError) as err: self.report_error('postprocessing: %s' % str(err)) return + try: + for ph in self._post_hooks: + ph(filename) + except Exception as err: + self.report_error('post hooks: %s' % str(err)) + return must_record_download_archive = True if must_record_download_archive or self.params.get('force_write_download_archive', False):