[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

@@ -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

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):
     """

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

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:
         """

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):
         """

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,))

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
 

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
 

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

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)