From c574be85f10de80482fa22a05d0b95bcc53a867c Mon Sep 17 00:00:00 2001 From: Kieran Eglin Date: Tue, 23 Apr 2024 10:22:36 -0700 Subject: [PATCH] Refactored MoveFilesPP to respect non-video files --- yt_dlp/YoutubeDL.py | 55 ++++++++--- yt_dlp/postprocessor/ffmpeg.py | 25 +++-- .../postprocessor/movefilesafterdownload.py | 93 +++++++++++++------ 3 files changed, 125 insertions(+), 48 deletions(-) diff --git a/yt_dlp/YoutubeDL.py b/yt_dlp/YoutubeDL.py index 9f730d038..5f652c637 100644 --- a/yt_dlp/YoutubeDL.py +++ b/yt_dlp/YoutubeDL.py @@ -3236,13 +3236,15 @@ def check_max_downloads(): sub_files = self._write_subtitles(info_dict, temp_filename) if sub_files is None: return - files_to_move.update(dict(sub_files)) + + files_to_move['requested_subtitles'] = sub_files thumb_files = self._write_thumbnails( 'video', info_dict, temp_filename, self.prepare_filename(info_dict, 'thumbnail')) if thumb_files is None: return - files_to_move.update(dict(thumb_files)) + + files_to_move['thumbnails'] = thumb_files infofn = self.prepare_filename(info_dict, 'infojson') _infojson_written = self._write_info_json('video', info_dict, infofn) @@ -3452,7 +3454,6 @@ def correct_ext(filename, ext=new_ext): dl_filename = dl_filename or temp_filename info_dict['__finaldir'] = os.path.dirname(os.path.abspath(encodeFilename(full_filename))) - except network_exceptions as err: self.report_error('unable to download video data: %s' % error_to_compat_str(err)) return @@ -3644,8 +3645,6 @@ def _delete_downloaded_files(self, *files_to_delete, info={}, msg=None): os.remove(filename) except OSError: self.report_warning(f'Unable to delete file {filename}') - if filename in info.get('__files_to_move', []): # NB: Delete even if None - del info['__files_to_move'][filename] @staticmethod def post_extract(info_dict): @@ -3675,10 +3674,7 @@ def run_pp(self, pp, infodict): if not files_to_delete: return infodict - if self.params.get('keepvideo', False): - for f in files_to_delete: - infodict['__files_to_move'].setdefault(f, '') - else: + if not self.params.get('keepvideo', False): self._delete_downloaded_files( *files_to_delete, info=infodict, msg='Deleting original file %s (pass -k to keep)') return infodict @@ -4294,10 +4290,17 @@ def _write_subtitles(self, info_dict, filename): sub_filename = subtitles_filename(filename, sub_lang, sub_format, info_dict.get('ext')) sub_filename_final = subtitles_filename(sub_filename_base, sub_lang, sub_format, info_dict.get('ext')) existing_sub = self.existing_file((sub_filename_final, sub_filename)) + if existing_sub: self.to_screen(f'[info] Video subtitle {sub_lang}.{sub_format} is already present') sub_info['filepath'] = existing_sub - ret.append((existing_sub, sub_filename_final)) + ret.append({ + 'current_filepath': existing_sub, + 'final_filepath': sub_filename_final, + 'lang': sub_lang, + 'ext': sub_info['ext'] + }) + continue self.to_screen(f'[info] Writing video subtitles to: {sub_filename}') @@ -4308,7 +4311,13 @@ def _write_subtitles(self, info_dict, filename): with open(sub_filename, 'w', encoding='utf-8', newline='') as subfile: subfile.write(sub_info['data']) sub_info['filepath'] = sub_filename - ret.append((sub_filename, sub_filename_final)) + ret.append({ + 'current_filepath': sub_filename, + 'final_filepath': sub_filename_final, + 'lang': sub_lang, + 'ext': sub_info['ext'] + }) + continue except OSError: self.report_error(f'Cannot write video subtitles file {sub_filename}') @@ -4319,7 +4328,13 @@ def _write_subtitles(self, info_dict, filename): sub_copy.setdefault('http_headers', info_dict.get('http_headers')) self.dl(sub_filename, sub_copy, subtitle=True) sub_info['filepath'] = sub_filename - ret.append((sub_filename, sub_filename_final)) + ret.append({ + 'current_filepath': sub_filename, + 'final_filepath': sub_filename_final, + 'lang': sub_lang, + 'ext': sub_info['ext'] + }) + except (DownloadError, ExtractorError, IOError, OSError, ValueError) + network_exceptions as err: msg = f'Unable to download video subtitles for {sub_lang!r}: {err}' if self.params.get('ignoreerrors') is not True: # False or 'only_download' @@ -4360,7 +4375,12 @@ def _write_thumbnails(self, label, info_dict, filename, thumb_filename_base=None self.to_screen('[info] %s is already present' % ( thumb_display_id if multiple else f'{label} thumbnail').capitalize()) t['filepath'] = existing_thumb - ret.append((existing_thumb, thumb_filename_final)) + ret.append({ + 'current_filepath': existing_thumb, + 'final_filepath': thumb_filename_final, + 'id': t['id'] + }) + else: self.to_screen(f'[info] Downloading {thumb_display_id} ...') try: @@ -4368,7 +4388,13 @@ def _write_thumbnails(self, label, info_dict, filename, thumb_filename_base=None self.to_screen(f'[info] Writing {thumb_display_id} to: {thumb_filename}') with open(encodeFilename(thumb_filename), 'wb') as thumbf: shutil.copyfileobj(uf, thumbf) - ret.append((thumb_filename, thumb_filename_final)) + + ret.append({ + 'current_filepath': thumb_filename, + 'final_filepath': thumb_filename_final, + 'id': t['id'] + }) + t['filepath'] = thumb_filename except network_exceptions as err: if isinstance(err, HTTPError) and err.status == 404: @@ -4378,4 +4404,5 @@ def _write_thumbnails(self, label, info_dict, filename, thumb_filename_base=None thumbnails.pop(idx) if ret and not write_all: break + return ret diff --git a/yt_dlp/postprocessor/ffmpeg.py b/yt_dlp/postprocessor/ffmpeg.py index 7d7f3f0eb..e794c1358 100644 --- a/yt_dlp/postprocessor/ffmpeg.py +++ b/yt_dlp/postprocessor/ffmpeg.py @@ -1016,8 +1016,10 @@ def run(self, info): 'filepath': new_file, } - info['__files_to_move'][new_file] = replace_extension( - info['__files_to_move'][sub['filepath']], new_ext) + for sub_info in info['__files_to_move']['requested_subtitles']: + if sub_info['lang'] == lang and sub_info['ext'] == sub['ext']: + sub_info['current_filepath'] = replace_extension(sub_info['current_filepath'], new_ext) + sub_info['final_filepath'] = replace_extension(sub_info['final_filepath'], new_ext) return sub_filenames, info @@ -1083,16 +1085,20 @@ def is_webp(cls, path): return imghdr.what(path) == 'webp' def fixup_webp(self, info, idx=-1): - thumbnail_filename = info['thumbnails'][idx]['filepath'] + thumbnail = info['thumbnails'][idx] + thumbnail_filename = thumbnail['filepath'] _, thumbnail_ext = os.path.splitext(thumbnail_filename) if thumbnail_ext: if thumbnail_ext.lower() != '.webp' and imghdr.what(thumbnail_filename) == 'webp': self.to_screen('Correcting thumbnail "%s" extension to webp' % thumbnail_filename) webp_filename = replace_extension(thumbnail_filename, 'webp') os.replace(thumbnail_filename, webp_filename) - info['thumbnails'][idx]['filepath'] = webp_filename - info['__files_to_move'][webp_filename] = replace_extension( - info['__files_to_move'].pop(thumbnail_filename), 'webp') + thumbnail['filepath'] = webp_filename + + for thumb_info in info['__files_to_move']['thumbnails']: + if thumb_info['id'] == thumbnail['id']: + thumb_info['current_filepath'] = replace_extension(thumbnail_filename, 'webp') + thumb_info['final_filepath'] = replace_extension(thumb_info['final_filepath'], 'webp') @staticmethod def _options(target_ext): @@ -1130,8 +1136,11 @@ def run(self, info): continue thumbnail_dict['filepath'] = self.convert_thumbnail(original_thumbnail, target_ext) files_to_delete.append(original_thumbnail) - info['__files_to_move'][thumbnail_dict['filepath']] = replace_extension( - info['__files_to_move'][original_thumbnail], target_ext) + + for thumb_info in info['__files_to_move']['thumbnails']: + if thumb_info['id'] == thumbnail_dict['id']: + thumb_info['current_filepath'] = replace_extension(thumb_info['current_filepath'], target_ext) + thumb_info['final_filepath'] = replace_extension(thumb_info['final_filepath'], target_ext) if not has_thumbnail: self.to_screen('There aren\'t any thumbnails to convert') diff --git a/yt_dlp/postprocessor/movefilesafterdownload.py b/yt_dlp/postprocessor/movefilesafterdownload.py index 23b09248c..d990eedf6 100644 --- a/yt_dlp/postprocessor/movefilesafterdownload.py +++ b/yt_dlp/postprocessor/movefilesafterdownload.py @@ -4,13 +4,11 @@ from ..compat import shutil from ..utils import ( PostProcessingError, - decodeFilename, - encodeFilename, make_dir, ) - class MoveFilesAfterDownloadPP(PostProcessor): + FILETYPE_KEYS = ['media', 'thumbnails', 'requested_subtitles'] def __init__(self, downloader=None, downloaded=True): PostProcessor.__init__(self, downloader) @@ -20,34 +18,77 @@ def __init__(self, downloader=None, downloaded=True): def pp_key(cls): return 'MoveFiles' + def expand_relative_paths(self, files_to_move, finaldir): + for filetype in self.FILETYPE_KEYS: + if filetype not in files_to_move: + continue + for file_attrs in files_to_move[filetype]: + if not os.path.isabs(file_attrs['final_filepath']): + file_attrs['final_filepath'] = os.path.join(finaldir, file_attrs['final_filepath']) + if not os.path.isabs(file_attrs['current_filepath']): + file_attrs['current_filepath'] = os.path.abspath(file_attrs['current_filepath']) + + return files_to_move + + def write_filepath_into_info(self, info, filetype, file_attrs): + if filetype == 'media': + info['filepath'] = file_attrs['final_filepath'] + return + + elif filetype == 'thumbnails': + for filetype_dict in info[filetype]: + if filetype_dict['id'] == file_attrs['id']: + filetype_dict['filepath'] = file_attrs['final_filepath'] + + elif filetype == 'requested_subtitles': + lang = file_attrs['lang'] + if lang in info[filetype]: + info[filetype][lang]['filepath'] = file_attrs['final_filepath'] + def run(self, info): - dl_path, dl_name = os.path.split(encodeFilename(info['filepath'])) + dl_path, dl_name = os.path.split(info['filepath']) finaldir = info.get('__finaldir', dl_path) - finalpath = os.path.join(finaldir, dl_name) + info['__files_to_move']['media'] = [] + if self._downloaded: - info['__files_to_move'][info['filepath']] = decodeFilename(finalpath) + info['__files_to_move']['media'] = [{ 'current_filepath': info['filepath'], 'final_filepath': dl_name }] - make_newfilename = lambda old: decodeFilename(os.path.join(finaldir, os.path.basename(encodeFilename(old)))) - for oldfile, newfile in info['__files_to_move'].items(): - if not newfile: - newfile = make_newfilename(oldfile) - if os.path.abspath(encodeFilename(oldfile)) == os.path.abspath(encodeFilename(newfile)): + files_to_move = self.expand_relative_paths(info['__files_to_move'], finaldir) + + for filetype in self.FILETYPE_KEYS: + if filetype not in files_to_move: continue - if not os.path.exists(encodeFilename(oldfile)): - self.report_warning('File "%s" cannot be found' % oldfile) - continue - if os.path.exists(encodeFilename(newfile)): - if self.get_param('overwrites', True): - self.report_warning('Replacing existing file "%s"' % newfile) - os.remove(encodeFilename(newfile)) - else: - self.report_warning( - 'Cannot move file "%s" out of temporary directory since "%s" already exists. ' - % (oldfile, newfile)) + + for file_attrs in files_to_move[filetype]: + current_filepath = file_attrs['current_filepath'] + final_filepath = file_attrs['final_filepath'] + + if not current_filepath or not final_filepath: continue - make_dir(newfile, PostProcessingError) - self.to_screen(f'Moving file "{oldfile}" to "{newfile}"') - shutil.move(oldfile, newfile) # os.rename cannot move between volumes - info['filepath'] = finalpath + if current_filepath == final_filepath: + # This ensures the infojson contains the full filepath even + # when --no-overwrites is used + self.write_filepath_into_info(info, filetype, file_attrs) + continue + + if not os.path.exists(current_filepath): + self.report_warning('File "%s" cannot be found' % current_filepath) + continue + + if os.path.exists(final_filepath): + if self.get_param('overwrites', True): + self.report_warning('Replacing existing file "%s"' % final_filepath) + os.remove(final_filepath) + else: + self.report_warning( + 'Cannot move file "%s" out of temporary directory since "%s" already exists. ' + % (current_filepath, final_filepath)) + continue + + make_dir(final_filepath, PostProcessingError) + self.to_screen(f'Moving file "{current_filepath}" to "{final_filepath}"') + shutil.move(current_filepath, final_filepath) # os.rename cannot move between volumes + self.write_filepath_into_info(info, filetype, file_attrs) + return [], info