15106d19fe108c42e87c83414faa666f9b7b167e
[panikdb.git] / panikdb / regie / management / commands / ytdl.py
1 import glob
2 import json
3 import os
4
5 import youtube_dl
6 from django.conf import settings
7 from django.core.files.storage import default_storage
8 from django.core.management.base import BaseCommand
9
10
11 class Command(BaseCommand):
12     def handle(self, verbosity, **kwargs):
13         self.verbosity = verbosity
14         self.ytdl_path = default_storage.path('ytdl')
15         if not os.path.exists(self.ytdl_path):
16             os.mkdir(self.ytdl_path)
17         while True:
18             done = []
19             for filename in glob.glob(os.path.join(self.ytdl_path, '*.json')):
20                 with open(filename) as fd:
21                     data = json.load(fd)
22
23                 if data.get('_status') == 'finished' or data.get('_error'):
24                     continue
25
26                 self.ytdl(filename, data)
27                 done.append(filename)
28             if not done:
29                 # nothing
30                 break
31
32     def ytdl(self, filename, data):
33         url = data['_url']
34         if self.verbosity:
35             print(url)
36
37         def progress_hook(d):
38             data['_percent_str'] = d.get('_percent_str')
39             data['_status'] = d.get('status')
40             with open(filename, 'w') as fd:
41                 json.dump(data, fp=fd, indent=2)
42
43         ydl_opts = {
44             'quiet': True,
45             'noplaylist': True,
46             'format': 'bestaudio',
47             'outtmpl': os.path.join(self.ytdl_path, '%(id)s.%(ext)s'),
48             'updatetime': False,
49             'postprocessors': [
50                 {
51                     'key': 'FFmpegExtractAudio',
52                     'preferredcodec': 'opus',
53                 }
54             ],
55             'progress_hooks': [progress_hook],
56         }
57         with youtube_dl.YoutubeDL(ydl_opts) as ydl:
58             ydl.download([url])
59
60         if data.get('_status') == 'finished':
61             base_filename = os.path.splitext(filename)[0]
62             real_filename = [
63                 x
64                 for x in glob.glob(base_filename + '*')
65                 if not x.endswith('.part') and not x.endswith('.json')
66             ][0]
67             url = os.path.join(settings.MEDIA_URL, 'ytdl', real_filename[len(self.ytdl_path) + 1 :])
68             data['_audio_url'] = url
69             with open(filename, 'w') as fd:
70                 json.dump(data, fp=fd, indent=2)
71         else:
72             data['_error'] = 'XXX'
73             with open(filename, 'w') as fd:
74                 json.dump(data, fp=fd, indent=2)