From 96e741c71d6313194bf923e231d4084443efc34e Mon Sep 17 00:00:00 2001 From: Jesse Bannon Date: Wed, 30 Nov 2022 10:36:13 -0800 Subject: [PATCH] [BACKEND] Proactively clean working directory (#379) * fix playlist last thumbnail move to copy * playlist working * channel working * working directory getting cleaned * working dir e2e fixture * cleanup files in any case * clean up if pre-split returns None --- src/ytdl_sub/downloaders/downloader.py | 5 +- src/ytdl_sub/plugins/nfo_tags.py | 3 + src/ytdl_sub/plugins/subtitles.py | 11 +++ .../subscriptions/subscription_download.py | 79 +++++++++++++------ src/ytdl_sub/utils/file_handler.py | 3 + src/ytdl_sub/utils/thumbnail.py | 2 + .../enhanced_download_archive.py | 11 +-- tests/conftest.py | 10 +-- tests/e2e/conftest.py | 78 ++++++++++++++---- tests/e2e/youtube/test_playlist.py | 4 +- 10 files changed, 152 insertions(+), 54 deletions(-) diff --git a/src/ytdl_sub/downloaders/downloader.py b/src/ytdl_sub/downloaders/downloader.py index 4456cdbc..84fa9a6f 100644 --- a/src/ytdl_sub/downloaders/downloader.py +++ b/src/ytdl_sub/downloaders/downloader.py @@ -422,7 +422,7 @@ class Downloader(DownloadArchiver, Generic[DownloaderOptionsT], ABC): FileHandler.delete(file_path=archive_path) # If the archive file did exist, restore the backup elif archive_file_exists: - FileHandler.copy(src_file_path=backup_archive_path, dst_file_path=archive_path) + FileHandler.move(src_file_path=backup_archive_path, dst_file_path=archive_path) # Clear info json files if true if clear_info_json_files: @@ -579,9 +579,9 @@ class Downloader(DownloadArchiver, Generic[DownloaderOptionsT], ABC): "Beginning downloads for %s", self.overrides.apply_formatter(collection_url.url) ) for entry in self._download(parents=parents, orphans=orphan_entries): - yield entry # Update thumbnails in case of last_entry self._download_url_thumbnails(collection_url=collection_url, entry=entry) + yield entry @classmethod def _download_thumbnail( @@ -630,6 +630,7 @@ class Downloader(DownloadArchiver, Generic[DownloaderOptionsT], ABC): self.save_file( file_name=entry.get_download_thumbnail_name(), output_file_name=thumbnail_name, + copy_file=True, ) self._url_state.thumbnails_downloaded.add(thumbnail_name) continue diff --git a/src/ytdl_sub/plugins/nfo_tags.py b/src/ytdl_sub/plugins/nfo_tags.py index ffd8c225..4bde7d4c 100644 --- a/src/ytdl_sub/plugins/nfo_tags.py +++ b/src/ytdl_sub/plugins/nfo_tags.py @@ -10,6 +10,7 @@ from typing import Optional from ytdl_sub.entries.entry import Entry from ytdl_sub.plugins.plugin import Plugin from ytdl_sub.plugins.plugin import PluginOptions +from ytdl_sub.utils.file_handler import FileHandler from ytdl_sub.utils.file_handler import FileMetadata from ytdl_sub.utils.xml import XmlElement from ytdl_sub.utils.xml import to_max_3_byte_utf8_dict @@ -171,6 +172,8 @@ class SharedNfoTagsPlugin(Plugin[SharedNfoTagsOptions], ABC): else: self.save_file(file_name=nfo_file_name, file_metadata=nfo_metadata) + FileHandler.delete(nfo_file_path) + class NfoTagsOptions(SharedNfoTagsOptions): """ diff --git a/src/ytdl_sub/plugins/subtitles.py b/src/ytdl_sub/plugins/subtitles.py index de98f888..8fefa973 100644 --- a/src/ytdl_sub/plugins/subtitles.py +++ b/src/ytdl_sub/plugins/subtitles.py @@ -1,3 +1,4 @@ +from pathlib import Path from typing import Dict from typing import List from typing import Optional @@ -7,6 +8,7 @@ from ytdl_sub.downloaders.ytdl_options_builder import YTDLOptionsBuilder from ytdl_sub.entries.entry import Entry from ytdl_sub.plugins.plugin import Plugin from ytdl_sub.plugins.plugin import PluginOptions +from ytdl_sub.utils.file_handler import FileHandler from ytdl_sub.utils.file_handler import FileMetadata from ytdl_sub.utils.logger import Logger from ytdl_sub.validators.string_formatter_validators import StringFormatterValidator @@ -206,4 +208,13 @@ class SubtitlesPlugin(Plugin[SubtitleOptions]): entry=entry, ) + # Delete any possible original subtitle files before conversion + # Can happen for both file and embedded subs + for lang in langs: + for possible_ext in SUBTITLE_EXTENSIONS: + possible_subs_file = ( + Path(self.working_directory) / f"{entry.uid}.{lang}.{possible_ext}" + ) + FileHandler.delete(possible_subs_file) + return file_metadata diff --git a/src/ytdl_sub/subscriptions/subscription_download.py b/src/ytdl_sub/subscriptions/subscription_download.py index 380fd155..7cdc2144 100644 --- a/src/ytdl_sub/subscriptions/subscription_download.py +++ b/src/ytdl_sub/subscriptions/subscription_download.py @@ -11,6 +11,7 @@ from ytdl_sub.subscriptions.base_subscription import BaseSubscription from ytdl_sub.subscriptions.subscription_ytdl_options import SubscriptionYTDLOptions from ytdl_sub.utils.datetime import to_date_range from ytdl_sub.utils.exceptions import ValidationException +from ytdl_sub.utils.file_handler import FileHandler from ytdl_sub.utils.file_handler import FileHandlerTransactionLog from ytdl_sub.utils.file_handler import FileMetadata from ytdl_sub.utils.thumbnail import convert_download_thumbnail @@ -94,6 +95,10 @@ class SubscriptionDownload(BaseSubscription, ABC): entry=entry, ) + def _delete_working_directory(self, is_error: bool) -> None: + _ = is_error + shutil.rmtree(self.working_directory) + @contextlib.contextmanager def _prepare_working_directory(self): """ @@ -104,8 +109,11 @@ class SubscriptionDownload(BaseSubscription, ABC): try: yield - finally: - shutil.rmtree(self.working_directory) + except Exception as exc: + self._delete_working_directory(is_error=True) + raise exc + else: + self._delete_working_directory(is_error=False) @contextlib.contextmanager def _maintain_archive_file(self): @@ -129,6 +137,8 @@ class SubscriptionDownload(BaseSubscription, ABC): self._enhanced_download_archive.remove_stale_files(date_range=date_range_to_keep) self._enhanced_download_archive.save_download_mappings() + FileHandler.delete(self._enhanced_download_archive.archive_working_file_path) + FileHandler.delete(self._enhanced_download_archive.mapping_working_file_path) @contextlib.contextmanager def _subscription_download_context_managers(self) -> None: @@ -156,6 +166,12 @@ class SubscriptionDownload(BaseSubscription, ABC): return plugins + @classmethod + def _cleanup_entry_files(cls, entry: Entry): + FileHandler.delete(entry.get_download_file_path()) + FileHandler.delete(entry.get_download_thumbnail_path()) + FileHandler.delete(entry.get_download_info_json_path()) + def _post_process_entry( self, plugins: List[Plugin], dry_run: bool, entry: Entry, entry_metadata: FileMetadata ): @@ -177,19 +193,26 @@ class SubscriptionDownload(BaseSubscription, ABC): def _process_entry( self, plugins: List[Plugin], dry_run: bool, entry: Entry, entry_metadata: FileMetadata ) -> None: + entry_: Optional[Entry] = entry + # First, modify the entry with all plugins for plugin in sorted(plugins, key=lambda _plugin: _plugin.priority.modify_entry): - # Return if it is None, it is indicated to not process any further - if (entry := plugin.modify_entry(entry)) is None: - return + # Break if it is None, it is indicated to not process any further + if (entry_ := plugin.modify_entry(entry_)) is None: + break - self._post_process_entry( - plugins=plugins, dry_run=dry_run, entry=entry, entry_metadata=entry_metadata - ) + if entry_: + self._post_process_entry( + plugins=plugins, dry_run=dry_run, entry=entry_, entry_metadata=entry_metadata + ) + + self._cleanup_entry_files(entry) def _process_split_entry( self, split_plugin: Plugin, plugins: List[Plugin], dry_run: bool, entry: Entry ) -> None: + entry_: Optional[Entry] = entry + plugins_pre_split = sorted( [plugin for plugin in plugins if not plugin.priority.modify_entry_after_split], key=lambda _plugin: _plugin.priority.modify_entry, @@ -202,27 +225,33 @@ class SubscriptionDownload(BaseSubscription, ABC): # First, modify the entry with pre_split plugins for plugin in plugins_pre_split: - # Return if it is None, it is indicated to not process any further - if (entry := plugin.modify_entry(entry)) is None: - return + # Break if it is None, it is indicated to not process any further + if (entry_ := plugin.modify_entry(entry_)) is None: + break # Then, perform the split - for split_entry, split_entry_metadata in split_plugin.split(entry=entry): + if entry_: + for split_entry, split_entry_metadata in split_plugin.split(entry=entry_): + split_entry_: Optional[Entry] = split_entry - for plugin in plugins_post_split: - # Return if it is None, it is indicated to not process any further. - # Break out of the plugin loop - if (split_entry := plugin.modify_entry(split_entry)) is None: - break + for plugin in plugins_post_split: + # Return if it is None, it is indicated to not process any further. + # Break out of the plugin loop + if (split_entry_ := plugin.modify_entry(split_entry_)) is None: + break - # If split_entry is None from modify_entry, do not post process - if split_entry: - self._post_process_entry( - plugins=plugins, - dry_run=dry_run, - entry=split_entry, - entry_metadata=split_entry_metadata, - ) + # If split_entry is None from modify_entry, do not post process + if split_entry_: + self._post_process_entry( + plugins=plugins, + dry_run=dry_run, + entry=split_entry_, + entry_metadata=split_entry_metadata, + ) + + self._cleanup_entry_files(split_entry) + + self._cleanup_entry_files(entry) def download(self, dry_run: bool = False) -> FileHandlerTransactionLog: """ diff --git a/src/ytdl_sub/utils/file_handler.py b/src/ytdl_sub/utils/file_handler.py index 2b8f1ced..ddabadd4 100644 --- a/src/ytdl_sub/utils/file_handler.py +++ b/src/ytdl_sub/utils/file_handler.py @@ -414,6 +414,9 @@ class FileHandler: self.copy(src_file_path=source_file_path, dst_file_path=output_file_path) else: self.move(src_file_path=source_file_path, dst_file_path=output_file_path) + # Simulate the file being moved during dry run by deleting it + elif self.dry_run and not copy_file: + FileHandler.delete(source_file_path) def delete_file_from_output_directory(self, file_name: str): """ diff --git a/src/ytdl_sub/utils/thumbnail.py b/src/ytdl_sub/utils/thumbnail.py index 25a4b9f5..2b3ee102 100644 --- a/src/ytdl_sub/utils/thumbnail.py +++ b/src/ytdl_sub/utils/thumbnail.py @@ -5,6 +5,7 @@ from urllib.request import urlopen from ytdl_sub.entries.entry import Entry from ytdl_sub.utils.ffmpeg import FFMPEG +from ytdl_sub.utils.file_handler import FileHandler from ytdl_sub.utils.retry import retry @@ -42,6 +43,7 @@ def convert_download_thumbnail(entry: Entry, error_if_not_found: bool = True) -> if not download_thumbnail_path == download_thumbnail_path_as_jpg: FFMPEG.run(["-bitexact", "-i", download_thumbnail_path, download_thumbnail_path_as_jpg]) + FileHandler.delete(download_thumbnail_path) @retry(times=5, exceptions=(Exception,)) diff --git a/src/ytdl_sub/ytdl_additions/enhanced_download_archive.py b/src/ytdl_sub/ytdl_additions/enhanced_download_archive.py index cbb3b211..a5e1ce37 100644 --- a/src/ytdl_sub/ytdl_additions/enhanced_download_archive.py +++ b/src/ytdl_sub/ytdl_additions/enhanced_download_archive.py @@ -400,8 +400,9 @@ class EnhancedDownloadArchive: Returns ------- The download archive's file name (no path) + Used to recreate yt-dlp's download archive in the working directory """ - return f".ytdl-subscribe-{self.subscription_name}-download-archive.txt" + return f".ytdl-sub-{self.subscription_name}-download-archive.txt" @property def working_directory(self) -> str: @@ -440,7 +441,7 @@ class EnhancedDownloadArchive: return str(Path(self.output_directory) / self._mapping_file_name) @property - def _mapping_working_file_path(self) -> str: + def mapping_working_file_path(self) -> str: """ Returns ------- @@ -449,7 +450,7 @@ class EnhancedDownloadArchive: return str(Path(self.working_directory) / self._mapping_file_name) @property - def _archive_working_file_path(self) -> str: + def archive_working_file_path(self) -> str: """ Returns ------- @@ -503,7 +504,7 @@ class EnhancedDownloadArchive: return self # Otherwise, create a ytdl download archive file in the working directory. - self.mapping.to_download_archive().to_file(self._archive_working_file_path) + self.mapping.to_download_archive().to_file(self.archive_working_file_path) return self @@ -555,7 +556,7 @@ class EnhancedDownloadArchive: self """ if not self.get_file_handler_transaction_log().is_empty: - self._download_mapping.to_file(output_json_file=self._mapping_working_file_path) + self._download_mapping.to_file(output_json_file=self.mapping_working_file_path) self.save_file_to_output_directory(file_name=self._mapping_file_name) return self diff --git a/tests/conftest.py b/tests/conftest.py index 7e863503..356e0d52 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,9 +1,7 @@ import contextlib import json import logging -import shutil import tempfile -from pathlib import Path from typing import Any from typing import Callable from typing import Dict @@ -15,14 +13,14 @@ import pytest from ytdl_sub.utils.logger import Logger -@pytest.fixture() -def output_directory(): +@pytest.fixture +def working_directory() -> str: with tempfile.TemporaryDirectory() as temp_dir: yield temp_dir -@pytest.fixture -def working_directory() -> str: +@pytest.fixture() +def output_directory(): with tempfile.TemporaryDirectory() as temp_dir: yield temp_dir diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py index b12160e8..de447ec1 100644 --- a/tests/e2e/conftest.py +++ b/tests/e2e/conftest.py @@ -1,3 +1,5 @@ +import json +import shutil import sys import tempfile from typing import List @@ -5,11 +7,42 @@ from typing import Tuple from unittest.mock import patch import pytest +from expected_download import _get_files_in_directory from ytdl_sub.cli.main import main from ytdl_sub.config.config_file import ConfigFile from ytdl_sub.subscriptions.subscription import Subscription +from ytdl_sub.subscriptions.subscription_download import SubscriptionDownload from ytdl_sub.utils.file_handler import FileHandlerTransactionLog +from ytdl_sub.utils.logger import Logger +from ytdl_sub.utils.yaml import load_yaml + +logger = Logger.get("test") + + +@pytest.fixture +def working_directory() -> str: + """ + Any time the working directory is used, ensure no files remain on cleaning it up + """ + with tempfile.TemporaryDirectory() as temp_dir: + + def _assert_working_directory_empty(self, is_error: bool): + files = [str(file_path) for file_path in _get_files_in_directory(temp_dir)] + num_files = len(files) + shutil.rmtree(temp_dir) + + if not is_error: + if num_files > 0: + logger.error("left-over files in working dir:\n%s", "\n".join(files)) + assert num_files == 0 + + with patch.object( + SubscriptionDownload, + "_delete_working_directory", + new=_assert_working_directory_empty, + ): + yield temp_dir @pytest.fixture() @@ -17,29 +50,46 @@ def music_video_config_path(): return "examples/music_videos_config.yaml" -@pytest.fixture() -def music_video_config(music_video_config_path): - return ConfigFile.from_file_path(config_path=music_video_config_path) +def _load_config(config_path: str, working_directory: str) -> ConfigFile: + config_dict = load_yaml(file_path=config_path) + config_dict["configuration"]["working_directory"] = working_directory + + return ConfigFile.from_dict(config_dict) @pytest.fixture() -def channel_as_tv_show_config(): - return ConfigFile.from_file_path(config_path="examples/tv_show_config.yaml") +def music_video_config(music_video_config_path, working_directory) -> ConfigFile: + return _load_config(music_video_config_path, working_directory) + + +@pytest.fixture() +def music_video_config_for_cli(music_video_config) -> str: + with tempfile.NamedTemporaryFile(suffix=".yaml") as tmp_file: + tmp_file.write(json.dumps(music_video_config._value).encode("utf-8")) + tmp_file.flush() + yield tmp_file.name + + +@pytest.fixture() +def channel_as_tv_show_config(working_directory) -> ConfigFile: + return _load_config( + config_path="examples/tv_show_config.yaml", working_directory=working_directory + ) @pytest.fixture -def soundcloud_discography_config(): - return ConfigFile.from_file_path(config_path="examples/soundcloud_discography_config.yaml") +def soundcloud_discography_config(working_directory) -> ConfigFile: + return _load_config( + config_path="examples/soundcloud_discography_config.yaml", + working_directory=working_directory, + ) @pytest.fixture() -def youtube_audio_config_path(): - return "examples/music_audio_from_videos.yaml" - - -@pytest.fixture() -def youtube_audio_config(youtube_audio_config_path): - return ConfigFile.from_file_path(config_path=youtube_audio_config_path) +def youtube_audio_config(working_directory) -> ConfigFile: + return _load_config( + config_path="examples/music_audio_from_videos.yaml", working_directory=working_directory + ) @pytest.fixture diff --git a/tests/e2e/youtube/test_playlist.py b/tests/e2e/youtube/test_playlist.py index 86265bd0..658d5872 100644 --- a/tests/e2e/youtube/test_playlist.py +++ b/tests/e2e/youtube/test_playlist.py @@ -97,7 +97,7 @@ class TestPlaylist: def test_playlist_download_from_cli_sub( self, preset_dict_to_subscription_yaml_generator, - music_video_config_path, + music_video_config_for_cli, playlist_preset_dict, output_directory, dry_run, @@ -106,7 +106,7 @@ class TestPlaylist: subscription_name="music_video_playlist_test", preset_dict=playlist_preset_dict ) as subscription_path: args = "--dry-run " if dry_run else "" - args += f"--config {music_video_config_path} " + args += f"--config {music_video_config_for_cli} " args += f"sub {subscription_path}" subscription_transaction_log = mock_run_from_cli(args=args)