Local playlists: save video thumbnails locally

1 files changed, 56 insertions, 3 deletions

diff --git a/youtube/local_playlist.py b/youtube/local_playlist.py
index 18857e5..76e7ee6 100644
--- a/youtube/local_playlist.py
+++ b/youtube/local_playlist.py
@@ -3,8 +3,11 @@ import json
 from youtube.template import Template
 from youtube import common
 import html
+import gevent
+import urllib
 
 playlists_directory = os.path.normpath("data/playlists")
+thumbnails_directory = os.path.normpath("data/playlist_thumbnails")
 with open('yt_local_playlist_template.html', 'r', encoding='utf-8') as file:
     local_playlist_template = Template(file.read())
 
@@ -20,13 +23,49 @@ def add_to_playlist(name, video_info_list):
     if not os.path.exists(playlists_directory):
         os.makedirs(playlists_directory)
     ids = video_ids_in_playlist(name)
+    missing_thumbnails = []
     with open(os.path.join(playlists_directory, name + ".txt"), "a", encoding='utf-8') as file:
         for info in video_info_list:
-            if json.loads(info)['id'] not in ids:
+            id = json.loads(info)['id']
+            if id not in ids:
                 file.write(info + "\n")
+                missing_thumbnails.append(id)
+    gevent.spawn(download_thumbnails, name, missing_thumbnails)
+
+def download_thumbnail(playlist_name, video_id):
+    url = "https://i.ytimg.com/vi/" + video_id + "/mqdefault.jpg"
+    save_location = os.path.join(thumbnails_directory, playlist_name, video_id + ".jpg")
+    try:
+        thumbnail = common.fetch_url(url, report_text="Saved local playlist thumbnail: " + video_id)
+    except urllib.error.HTTPError as e:
+        print("Failed to download thumbnail for " + video_id + ": " + str(e))
+        return
+    try:
+        f = open(save_location, 'wb')
+    except FileNotFoundError:
+        os.makedirs(os.path.join(thumbnails_directory, playlist_name))
+        f = open(save_location, 'wb')
+    f.write(thumbnail)
+    f.close()
+
+def download_thumbnails(playlist_name, ids):
+    # only do 5 at a time
+    # do the n where n is divisible by 5
+    i = -1
+    for i in range(0, int(len(ids)/5) - 1 ):
+        gevent.joinall([gevent.spawn(download_thumbnail, playlist_name, ids[j]) for j in range(i*5, i*5 + 5)])
+    # do the remainders (< 5)
+    gevent.joinall([gevent.spawn(download_thumbnail, playlist_name, ids[j]) for j in range(i*5 + 5, len(ids))])
+            
         
-        
+
 def get_local_playlist_page(name):
+    try:
+        thumbnails = set(os.listdir(os.path.join(thumbnails_directory, name)))
+    except FileNotFoundError:
+        thumbnails = set()
+    missing_thumbnails = []
+
     videos_html = ''
     with open(os.path.join(playlists_directory, name + ".txt"), 'r', encoding='utf-8') as file:
         videos = file.read()
@@ -34,10 +73,15 @@ def get_local_playlist_page(name):
     for video in videos:
         try:
             info = json.loads(video)
-            info['thumbnail'] = common.get_thumbnail_url(info['id'])
+            if info['id'] + ".jpg" in thumbnails:
+                info['thumbnail'] = "/youtube.com/data/playlist_thumbnails/" + name + "/" + info['id'] + ".jpg"
+            else:
+                info['thumbnail'] = common.get_thumbnail_url(info['id'])
+                missing_thumbnails.append(info['id'])
             videos_html += common.video_item_html(info, common.small_video_item_template)
         except json.decoder.JSONDecodeError:
             pass
+    gevent.spawn(download_thumbnails, name, missing_thumbnails)
     return local_playlist_template.substitute(
         page_title = name + ' - Local playlist',
         header = common.get_header(),
@@ -68,6 +112,15 @@ def remove_from_playlist(name, video_info_list):
     with open(os.path.join(playlists_directory, name + ".txt"), 'w', encoding='utf-8') as file:
         file.write("\n".join(videos_out) + "\n")
 
+    try:
+        thumbnails = set(os.listdir(os.path.join(thumbnails_directory, name)))
+    except FileNotFoundError:
+        pass
+    else:
+        to_delete = thumbnails & set(id + ".jpg" for id in ids)
+        for file in to_delete:
+            os.remove(os.path.join(thumbnails_directory, name, file))
+
 def get_playlists_list_page():
     page = '''<ul>\n'''
     list_item_template = Template('''    <li><a href="$url">$name</a></li>\n''')