Add post_hooks option to YoutubeDL.py (https://github.com/ytdl-org/youtube-dl/pull/27573)

Authored by: alexmerkel

5 files changed, 89 insertions, 3 deletions

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):