Merge pull request #32 from sutro-planet/thumbnail_api

Add an API to download yt video thumbnail. 增添视频缩略图接口

app.py

@@ -1,5 +1,6 @@
 import os
 import json
+import requests
 
 from pytubefix import YouTube
 from moviepy.editor import VideoFileClip
@@ -127,6 +128,34 @@ def video_upload():
         return jsonify({"message", log_error_return_str(f"Video upload failed: {file.filename} extension not allowed")})
 
 
+@app.route('/yt_thumbnail', methods=['POST'])
+@pytvzhen_api_request_counter
+@require_video_id_from_post_request
+def yt_thumbnail(video_id):
+    output_path = app.config['OUTPUT_PATH']
+    thumbnail_fn = f"{video_id}_thumbnail.png"
+
+    if os.path.isfile(thumbnail_fn):
+        return send_from_directory(output_path, thumbnail_fn, mimetype='image/png')
+
+    thumbnail_save_path = os.path.join(output_path, thumbnail_fn)
+    try:
+        yt = YouTube(f'https://www.youtube.com/watch?v={video_id}', proxies=None)
+
+        response = requests.get(yt.thumbnail_url)
+        if response.status_code == 200:
+            with open(thumbnail_save_path, 'wb') as file:
+                file.write(response.content)
+            return send_from_directory(output_path, thumbnail_fn, mimetype='image/png')
+
+        raise Exception(f"thumbnail download failed: {response.status_code} {response.content}")
+    except Exception as e:
+        exception = e
+
+    return jsonify({"message": log_error_return_str(
+        f'An error occurred while downloading video thumbnail {video_id} to {thumbnail_save_path}: {exception}')}), 500
+
+
 @app.route('/yt_download', methods=['POST'])
 @pytvzhen_api_request_counter
 @require_video_id_from_post_request
@@ -169,7 +198,7 @@ def yt_download(video_id):
         exception = e
 
     return jsonify({"message": log_error_return_str(
-        f'An error occurred while downloading video  {video_id} to {video_save_path}: {exception}')}), 500
+        f'An error occurred while downloading video {video_id} to {video_save_path}: {exception}')}), 500
 
 
 @app.route('/yt/<video_id>', methods=['GET'])

test_app.py

@@ -1,6 +1,7 @@
 import os
 import unittest
 import tempfile
+import hashlib
 
 from app import app
 from app import url_rule_to_base
@@ -48,6 +49,19 @@ def test_download_yt_video_pytube_issue(self):
         assert response.status_code != 500, "Pytube fix not applied, see https://github.com/pytube/pytube/pull/1312"
         assert response.status_code == 200
 
+    def test_download_yt_thumbnail(self):
+        response = self.app.post("/yt_thumbnail", json={'video_id': 'SrvXsYxbgC4'})
+        expected_fn = os.path.join(self.test_dir, 'SrvXsYxbgC4_thumbnail.png')
+
+        assert response.status_code == 200
+        assert os.path.isfile(expected_fn), "thumbnail not downloaded and cached on server side"
+
+        returned_sha = hashlib.sha256(response.data).hexdigest()
+
+        with open(expected_fn, 'rb', buffering=0) as f:
+            expected_sha = hashlib.sha256(f.read()).hexdigest()
+            assert returned_sha == expected_sha, f"thumbnail doesn't match, expected {expected_sha} actual {returned_sha}"
+
     def tearDown(self):
         for root, dirs, files in os.walk(self.test_dir, topdown=False):
             for name in files: