From 44bb6c205620d133a4772fb6de56786859aecfba Mon Sep 17 00:00:00 2001 From: Kieran Eglin Date: Wed, 24 Apr 2024 09:47:58 -0700 Subject: [PATCH] [WIP] got refactor of file mover basically working --- yt_dlp/YoutubeDL.py | 28 ++- .../postprocessor/movefilesafterdownload.py | 227 +++++++++++++----- yt_dlp/utils/_utils.py | 4 +- 3 files changed, 188 insertions(+), 71 deletions(-) diff --git a/yt_dlp/YoutubeDL.py b/yt_dlp/YoutubeDL.py index 5f652c637..ea4e4fe08 100644 --- a/yt_dlp/YoutubeDL.py +++ b/yt_dlp/YoutubeDL.py @@ -1982,8 +1982,9 @@ class YoutubeDL: 'playlist', ie_result, self.prepare_filename(ie_copy, 'pl_infojson')) if _infojson_written is None: return - if self._write_description('playlist', ie_result, - self.prepare_filename(ie_copy, 'pl_description')) is None: + + description_file = self._write_description('playlist', ie_result, self.prepare_filename(ie_copy, 'pl_description')) + if description_file is None: return # TODO: This should be passed to ThumbnailsConvertor if necessary self._write_thumbnails('playlist', ie_result, self.prepare_filename(ie_copy, 'pl_thumbnail')) @@ -3229,8 +3230,8 @@ class YoutubeDL: if not self._ensure_dir_exists(encodeFilename(temp_filename)): return - if self._write_description('video', info_dict, - self.prepare_filename(info_dict, 'description')) is None: + description_file = self._write_description('video', info_dict, self.prepare_filename(info_dict, 'description')) + if description_file is None: return sub_files = self._write_subtitles(info_dict, temp_filename) @@ -4245,27 +4246,28 @@ class YoutubeDL: self.report_error(f'Cannot write {label} metadata to JSON file {infofn}') return None - def _write_description(self, label, ie_result, descfn): + def _write_description(self, label, info_dict, filename): ''' Write description and returns True = written, False = skip, None = error ''' if not self.params.get('writedescription'): return False - elif not descfn: + elif not filename: self.write_debug(f'Skipping writing {label} description') return False - elif not self._ensure_dir_exists(descfn): + elif not self._ensure_dir_exists(filename): return None - elif not self.params.get('overwrites', True) and os.path.exists(descfn): + elif not self.params.get('overwrites', True) and os.path.exists(filename): self.to_screen(f'[info] {label.title()} description is already present') - elif ie_result.get('description') is None: + elif info_dict.get('description') is None: self.to_screen(f'[info] There\'s no {label} description to write') return False else: try: - self.to_screen(f'[info] Writing {label} description to: {descfn}') - with open(encodeFilename(descfn), 'w', encoding='utf-8') as descfile: - descfile.write(ie_result['description']) + self.to_screen(f'[info] Writing {label} description to: {filename}') + with open(filename, 'w', encoding='utf-8') as descfile: + descfile.write(info_dict['description']) + info_dict['description_filepath'] = filename except OSError: - self.report_error(f'Cannot write {label} description file {descfn}') + self.report_error(f'Cannot write {label} description file {filename}') return None return True diff --git a/yt_dlp/postprocessor/movefilesafterdownload.py b/yt_dlp/postprocessor/movefilesafterdownload.py index cb4c1cd1c..b80ba6898 100644 --- a/yt_dlp/postprocessor/movefilesafterdownload.py +++ b/yt_dlp/postprocessor/movefilesafterdownload.py @@ -1,15 +1,23 @@ import os +from pathlib import Path from .common import PostProcessor from ..compat import shutil from ..utils import ( PostProcessingError, make_dir, + replace_extension ) - +import pdb class MoveFilesAfterDownloadPP(PostProcessor): - FILETYPE_KEYS = ['media', 'thumbnails', 'requested_subtitles'] + TOP_LEVEL_KEYS = ['filepath'] + # Map of the keys that contain moveable files and the 'type' of the file + # for generating the output filename + CHILD_KEYS = { + 'thumbnails': 'thumbnail', + 'requested_subtitles': 'subtitle' + } def __init__(self, downloader=None, downloaded=True): PostProcessor.__init__(self, downloader) @@ -18,74 +26,179 @@ class MoveFilesAfterDownloadPP(PostProcessor): @classmethod def pp_key(cls): return 'MoveFiles' + + def move_file_and_write_to_info(self, info_dict, relevant_dict, output_file_type): + if 'filepath' not in relevant_dict: + return - def expand_relative_paths(self, files_to_move, finaldir): - for filetype in self.FILETYPE_KEYS: - if filetype not in files_to_move: - files_to_move[filetype] = [] + output_file_type = output_file_type or '' + current_filepath = relevant_dict['filepath'] + # This approach is needed to preserved indexed thumbnail paths from `--write-all-thumbnails` + # and also to support user-defined extensions (eg: `%(title)s.temp.%(ext)s`) + extension = ''.join(Path(current_filepath).suffixes) + name_format = self._downloader.prepare_filename(info_dict, output_file_type) + final_filepath = replace_extension(name_format, extension) + move_result = self.move_file(info_dict, current_filepath, final_filepath) + + print('*******************') + print("output_file_type", output_file_type) + print("name_format", name_format) + print("current_filepath", current_filepath) + print("final_filepath", final_filepath) + print("move_result", move_result) + print('*******************') - 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']) + if move_result: + relevant_dict['filepath'] = move_result + else: + del relevant_dict['filepath'] - return files_to_move + def move_file(self, info_dict, current_filepath, final_filepath): + if not current_filepath or not final_filepath: + return + + dl_path, _dl_name = os.path.split(info_dict['filepath']) + finaldir = info_dict.get('__finaldir', os.path.abspath(dl_path)) - def write_filepath_into_info(self, info, filetype, file_attrs): - if filetype == 'media': - info['filepath'] = file_attrs['final_filepath'] + if not os.path.isabs(current_filepath): + current_filepath = os.path.join(finaldir, current_filepath) - elif filetype == 'thumbnails': - for filetype_dict in info[filetype]: - if filetype_dict['id'] == file_attrs['id']: - filetype_dict['filepath'] = file_attrs['final_filepath'] + if not os.path.isabs(final_filepath): + final_filepath = os.path.join(finaldir, final_filepath) - elif filetype == 'requested_subtitles': - lang = file_attrs['lang'] - if lang in info[filetype]: - info[filetype][lang]['filepath'] = file_attrs['final_filepath'] + if current_filepath == final_filepath: + return final_filepath + + if not os.path.exists(current_filepath): + self.report_warning('File "%s" cannot be found' % current_filepath) + return + + 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)) + return + + 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 + + return final_filepath def run(self, info): dl_path, dl_name = os.path.split(info['filepath']) finaldir = info.get('__finaldir', os.path.abspath(dl_path)) + # TODO: add one single key to infodict with ALL downloaded files + # TODO: test with --path temp and stuff + # TODO: make the below work with not-currently-written filepaths like description, annotations, etc + # - Descriptions work, have to do all the other ones too + # - I lied, this should become another post-processor + # TODO: [DONE] probably something with relative paths into absolute again? + # TODO: remove all __files_to_move stuff when done + # TODO: add net-new filepaths to `sanitize_info` + # TODO: consider adding a `infojson_filepath` key in addition to the `infojson_filename` key where the former is the fullpath - if self._downloaded: - info['__files_to_move']['media'] = [{'current_filepath': info['filepath'], 'final_filepath': dl_name}] + for filepath_key in self.TOP_LEVEL_KEYS: + self.move_file_and_write_to_info(info, info, None) - files_to_move = self.expand_relative_paths(info['__files_to_move'], finaldir) + for key, output_file_type in self.CHILD_KEYS.items(): + if key not in info: + continue - for filetype in self.FILETYPE_KEYS: - for file_attrs in files_to_move[filetype]: - current_filepath = file_attrs['current_filepath'] - final_filepath = file_attrs['final_filepath'] + if isinstance(info[key], list) or isinstance(info[key], dict): + iterable = info[key].values() if isinstance(info[key], dict) else info[key] - if not current_filepath or not final_filepath: - continue - - 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) + for file_dict in iterable: + self.move_file_and_write_to_info(info, file_dict, output_file_type) return [], info + +# class MoveFilesAfterDownloadPP(PostProcessor): +# FILETYPE_KEYS = ['media', 'thumbnails', 'requested_subtitles'] + +# def __init__(self, downloader=None, downloaded=True): +# PostProcessor.__init__(self, downloader) +# self._downloaded = downloaded + +# @classmethod +# 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: +# files_to_move[filetype] = [] + +# 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'] + +# 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(info['filepath']) +# finaldir = info.get('__finaldir', os.path.abspath(dl_path)) + +# th = self._downloader.prepare_filename(info, 'thumbnail') +# pdb.set_trace() +# print("th", th) + +# if self._downloaded: +# info['__files_to_move']['media'] = [{'current_filepath': info['filepath'], 'final_filepath': dl_name}] + +# files_to_move = self.expand_relative_paths(info['__files_to_move'], finaldir) + +# for filetype in self.FILETYPE_KEYS: +# 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 + +# 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 diff --git a/yt_dlp/utils/_utils.py b/yt_dlp/utils/_utils.py index e3e80f3d3..cb018dbfd 100644 --- a/yt_dlp/utils/_utils.py +++ b/yt_dlp/utils/_utils.py @@ -2099,7 +2099,9 @@ def prepend_extension(filename, ext, expected_real_ext=None): def replace_extension(filename, ext, expected_real_ext=None): name, real_ext = os.path.splitext(filename) - return '{}.{}'.format( + ext = ext if ext.startswith('.') else '.' + ext + + return '{}{}'.format( name if not expected_real_ext or real_ext[1:] == expected_real_ext else filename, ext)