Authored by: alexmerkel
This commit is contained in:
Alex Merkel 2020-12-29 16:03:07 +01:00 committed by pukkandan
parent 62d80ba17c
commit ab8e5e516f
5 changed files with 89 additions and 3 deletions

View file

@ -50,7 +50,8 @@ offlinetest: codetest
--exclude test_subtitles.py \ --exclude test_subtitles.py \
--exclude test_write_annotations.py \ --exclude test_write_annotations.py \
--exclude test_youtube_lists.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 tar: youtube-dlc.tar.gz

View file

@ -1,7 +1,7 @@
@echo off @echo off
rem Keep this list in sync with the `offlinetest` target in Makefile 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" ( if "%YTDL_TEST_SET%" == "core" (
set test_set="-I test_("%DOWNLOAD_TESTS%")\.py" set test_set="-I test_("%DOWNLOAD_TESTS%")\.py"

View file

@ -1,7 +1,7 @@
#!/bin/bash #!/bin/bash
# Keep this list in sync with the `offlinetest` target in Makefile # 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="" test_set=""
multiprocess_args="" multiprocess_args=""

68
test/test_post_hooks.py Normal file
View file

@ -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()

View file

@ -252,6 +252,9 @@ class YoutubeDL(object):
youtube_dlc/postprocessor/__init__.py for a list. youtube_dlc/postprocessor/__init__.py for a list.
as well as any further keyword arguments for the as well as any further keyword arguments for the
postprocessor. 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_hooks: A list of functions that get called on download
progress, with a dictionary with the entries progress, with a dictionary with the entries
* status: One of "downloading", "error", or "finished". * status: One of "downloading", "error", or "finished".
@ -369,6 +372,7 @@ class YoutubeDL(object):
self._ies = [] self._ies = []
self._ies_instances = {} self._ies_instances = {}
self._pps = [] self._pps = []
self._post_hooks = []
self._progress_hooks = [] self._progress_hooks = []
self._download_retcode = 0 self._download_retcode = 0
self._num_downloads = 0 self._num_downloads = 0
@ -472,6 +476,9 @@ class YoutubeDL(object):
pp = pp_class(self, **compat_kwargs(pp_def)) pp = pp_class(self, **compat_kwargs(pp_def))
self.add_post_processor(pp) 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', []): for ph in self.params.get('progress_hooks', []):
self.add_progress_hook(ph) self.add_progress_hook(ph)
@ -524,6 +531,10 @@ class YoutubeDL(object):
self._pps.append(pp) self._pps.append(pp)
pp.set_downloader(self) 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): def add_progress_hook(self, ph):
"""Add the progress hook (currently only for the file downloader)""" """Add the progress hook (currently only for the file downloader)"""
self._progress_hooks.append(ph) self._progress_hooks.append(ph)
@ -2199,6 +2210,12 @@ class YoutubeDL(object):
except (PostProcessingError) as err: except (PostProcessingError) as err:
self.report_error('postprocessing: %s' % str(err)) self.report_error('postprocessing: %s' % str(err))
return 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 must_record_download_archive = True
if must_record_download_archive or self.params.get('force_write_download_archive', False): if must_record_download_archive or self.params.get('force_write_download_archive', False):