Use FILE_UPLOAD_TEMP_DIR for big files.
[audio.git] / src / youtube / utils.py
1 import os
2 import shutil
3 import subprocess
4 from tempfile import NamedTemporaryFile
5 from django.conf import settings
6
7
8 FILE_CACHE = getattr(settings, 'FILE_CACHE', 'file_cache/')
9
10
11 def link_or_copy(src, dst):
12     dstdir = os.path.dirname(dst)
13     if not os.path.exists(dstdir):
14         os.makedirs(dstdir)
15     if os.path.exists(dst):
16         os.unlink(dst)
17         # FIXME: tiny window here when the temp path is not taken.
18     try:
19         os.link(src, dst)
20     except OSError:
21         shutil.copyfile(src, dst)
22
23
24 def process_to_file(cmdline, prefix='', suffix='', cache_key=None, output_path=None):
25     if not output_path:
26         tmp = NamedTemporaryFile(
27             prefix=prefix, suffix=suffix, delete=False,
28             dir=settings.FILE_UPLOAD_TEMP_DIR
29         )
30         tmp.close()
31         output_path = tmp.name
32
33     if cache_key:
34         cache_path = FILE_CACHE + cache_key.replace('/', '__')
35
36     if cache_key and os.path.exists(cache_path):
37         link_or_copy(cache_path, output_path)
38     else:
39         # Actually run the processing.
40         subprocess.run(cmdline + [output_path], check=True)
41         if cache_key:
42             link_or_copy(output_path, cache_path)
43
44     return output_path
45
46
47 def video_from_image(img_path, duration, fps=25, cache=True):
48     return process_to_file(
49         ['ffmpeg', '-y', '-loop', '1', '-t', str(duration), '-i', img_path, '-c:v', 'libx264', '-vf', f'fps={fps},format=yuv420p'],
50         'image-',
51         '.mkv',
52         f'video_from_image:{img_path}:{duration}:{fps}.mkv' if cache else None
53     )
54
55
56 def cut_video(video_path, duration):
57     return process_to_file(
58         ['ffmpeg', '-y', '-i', video_path, '-t', str(duration), '-c', 'copy'],
59         'cut-',
60         '.mkv'
61     )
62
63
64 def ffmpeg_concat(paths, suffix, copy=False):
65     filelist = NamedTemporaryFile(
66         prefix='concat-', suffix='.txt',
67         dir=settings.FILE_UPLOAD_TEMP_DIR
68     )
69     for path in paths:
70         filelist.write(f"file '{path}'\n".encode('utf-8'))
71     filelist.flush()
72
73     args = ['ffmpeg', '-y', '-safe', '0', '-f', 'concat', '-i', filelist.name]
74     if copy:
75         args += ['-c', 'copy']
76     outname = process_to_file(args, 'concat-', suffix)
77
78     filelist.close()
79     return outname
80
81
82 def concat_videos(paths):
83     return ffmpeg_concat(paths, '.mkv', copy=True)
84
85
86 def concat_audio(paths):
87     return ffmpeg_concat(paths, '.flac')
88
89
90 def standardize_audio(p, cache=True):
91     return process_to_file(
92         ['ffmpeg', '-y', '-i', p, '-sample_fmt', 's16', '-acodec', 'flac', '-ac', '2', '-ar', '44100'],
93         'standardize-', '.flac',
94         f'standardize_audio:{p}.flac' if cache else None
95     )
96
97
98 def standardize_video(p, cache=True):
99     return process_to_file(
100         ['ffmpeg', '-y', '-i', p],
101         'standardize-', '.mkv',
102         f'standardize_video:{p}.mkv' if cache else None
103     )
104
105
106 def mux(channels, output_path=None):
107     args = ['ffmpeg', '-y']
108     for c in channels:
109         args.extend(['-i', c])
110     args.extend(['-c', 'copy'])
111     return process_to_file(args, 'mux-', '.mkv', output_path=output_path)
112
113
114 def get_duration(path):
115     return float(
116         subprocess.run(
117             [
118                 "ffprobe",
119                 "-i",
120                 path,
121                 "-show_entries",
122                 "format=duration",
123                 "-v",
124                 "quiet",
125                 "-of",
126                 "csv=p=0",
127             ],
128             capture_output=True,
129             text=True,
130             check=True,
131         ).stdout
132     )
133
134
135 def get_framerate(path):
136     rates = subprocess.run(
137             [
138                 "ffprobe",
139                 "-i",
140                 path,
141                 "-show_entries",
142                 "stream=r_frame_rate",
143                 "-v",
144                 "quiet",
145                 "-of",
146                 "csv=p=0",
147             ],
148             capture_output=True,
149             text=True,
150             check=True,
151         ).stdout.strip().split('\n')
152     for rate in rates:
153         a, b = rate.split('/')
154         if b == '1':
155             return int(a)