411 lines
16 KiB
Python
411 lines
16 KiB
Python
import contextlib
|
|
import json
|
|
import os
|
|
import shutil
|
|
from abc import ABC
|
|
from pathlib import Path
|
|
from typing import Dict
|
|
from typing import Iterable
|
|
from typing import List
|
|
from typing import Optional
|
|
from typing import Set
|
|
from typing import Tuple
|
|
|
|
from ytdl_sub.entries.entry import Entry
|
|
from ytdl_sub.plugins.plugin import Plugin
|
|
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.file_handler import get_file_extension
|
|
from ytdl_sub.utils.thumbnail import convert_download_thumbnail
|
|
from ytdl_sub.ytdl_additions.enhanced_download_archive import DownloadMappings
|
|
|
|
|
|
def _get_split_plugin(plugins: List[Plugin]) -> Optional[Plugin]:
|
|
split_plugins = [plugin for plugin in plugins if plugin.is_split_plugin]
|
|
|
|
if len(split_plugins) == 1:
|
|
return split_plugins[0]
|
|
if len(split_plugins) > 1:
|
|
raise ValidationException("Can not use more than one split plugins at a time")
|
|
return None
|
|
|
|
|
|
class SubscriptionDownload(BaseSubscription, ABC):
|
|
"""
|
|
Handles the subscription download logic
|
|
"""
|
|
|
|
def _move_entry_files_to_output_directory(
|
|
self,
|
|
dry_run: bool,
|
|
entry: Entry,
|
|
entry_metadata: Optional[FileMetadata] = None,
|
|
):
|
|
"""
|
|
Helper function to move the media file and optionally thumbnail file to the output directory
|
|
for a single entry.
|
|
|
|
Parameters
|
|
----------
|
|
dry_run
|
|
Whether this session is a dry-run or not
|
|
entry:
|
|
The entry with files to move
|
|
entry_metadata
|
|
Optional. Metadata to record to the transaction log for this entry
|
|
"""
|
|
# Move the file after all direct file modifications are complete
|
|
output_file_name = self.overrides.apply_formatter(
|
|
formatter=self.output_options.file_name, entry=entry
|
|
)
|
|
self._enhanced_download_archive.save_file_to_output_directory(
|
|
file_name=entry.get_download_file_name(),
|
|
file_metadata=entry_metadata,
|
|
output_file_name=output_file_name,
|
|
entry=entry,
|
|
)
|
|
|
|
# TODO: see if entry even has a thumbnail
|
|
if self.output_options.thumbnail_name:
|
|
output_thumbnail_name = self.overrides.apply_formatter(
|
|
formatter=self.output_options.thumbnail_name, entry=entry
|
|
)
|
|
|
|
# We always convert entry thumbnails to jpgs, and is performed here
|
|
if not dry_run:
|
|
convert_download_thumbnail(entry=entry)
|
|
|
|
# Copy the thumbnails since they could be used later for other things
|
|
self._enhanced_download_archive.save_file_to_output_directory(
|
|
file_name=entry.get_download_thumbnail_name(),
|
|
output_file_name=output_thumbnail_name,
|
|
entry=entry,
|
|
copy_file=True,
|
|
)
|
|
|
|
if self.output_options.info_json_name:
|
|
output_info_json_name = self.overrides.apply_formatter(
|
|
formatter=self.output_options.info_json_name, entry=entry
|
|
)
|
|
|
|
# if not dry-run, write the info json
|
|
if not dry_run:
|
|
entry.write_info_json()
|
|
|
|
self._enhanced_download_archive.save_file_to_output_directory(
|
|
file_name=entry.get_download_info_json_name(),
|
|
output_file_name=output_info_json_name,
|
|
entry=entry,
|
|
)
|
|
|
|
def _delete_working_directory(self, is_error: bool = False) -> None:
|
|
_ = is_error
|
|
if os.path.isdir(self.working_directory):
|
|
shutil.rmtree(self.working_directory)
|
|
|
|
@contextlib.contextmanager
|
|
def _prepare_working_directory(self):
|
|
"""
|
|
Context manager to create all directories to the working directory. Deletes the entire
|
|
working directory when cleaning up.
|
|
"""
|
|
self._delete_working_directory()
|
|
os.makedirs(self.working_directory, exist_ok=True)
|
|
|
|
try:
|
|
yield
|
|
except Exception as exc:
|
|
self._delete_working_directory(is_error=True)
|
|
raise exc
|
|
else:
|
|
self._delete_working_directory()
|
|
|
|
@contextlib.contextmanager
|
|
def _maintain_archive_file(self):
|
|
"""
|
|
Context manager to initialize the enhanced download archive
|
|
"""
|
|
if self.maintain_download_archive:
|
|
self._enhanced_download_archive.prepare_download_archive()
|
|
|
|
yield
|
|
|
|
# If output options maintains stale file deletion, perform the delete here prior to saving
|
|
# the download archive
|
|
if self.maintain_download_archive:
|
|
date_range_to_keep = to_date_range(
|
|
before=self.output_options.keep_files_before,
|
|
after=self.output_options.keep_files_after,
|
|
overrides=self.overrides,
|
|
)
|
|
if date_range_to_keep:
|
|
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 _remove_empty_directories_in_output_directory(self):
|
|
try:
|
|
yield
|
|
finally:
|
|
if not self._enhanced_download_archive.is_dry_run:
|
|
for root, dir_names, filenames in os.walk(Path(self.output_directory), topdown=False):
|
|
for dir_name in dir_names:
|
|
dir_path = Path(root) / dir_name
|
|
if len(os.listdir(dir_path)) == 0:
|
|
os.rmdir(dir_path)
|
|
|
|
@contextlib.contextmanager
|
|
def _subscription_download_context_managers(self) -> None:
|
|
with (
|
|
self._prepare_working_directory(),
|
|
self._maintain_archive_file(),
|
|
self._remove_empty_directories_in_output_directory(),
|
|
):
|
|
yield
|
|
|
|
def _initialize_plugins(self) -> List[Plugin]:
|
|
"""
|
|
Returns
|
|
-------
|
|
List of plugins defined in the subscription, initialized and ready to use.
|
|
"""
|
|
plugins: List[Plugin] = []
|
|
for plugin_type, plugin_options in self.plugins.zipped():
|
|
plugin = plugin_type(
|
|
plugin_options=plugin_options,
|
|
overrides=self.overrides,
|
|
enhanced_download_archive=self._enhanced_download_archive,
|
|
)
|
|
|
|
plugins.append(plugin)
|
|
|
|
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
|
|
):
|
|
# Post-process the entry with all plugins
|
|
for plugin in sorted(plugins, key=lambda _plugin: _plugin.priority.post_process):
|
|
optional_plugin_entry_metadata = plugin.post_process_entry(entry)
|
|
if optional_plugin_entry_metadata:
|
|
entry_metadata.extend(optional_plugin_entry_metadata)
|
|
|
|
# Then, move it to the output directory
|
|
self._move_entry_files_to_output_directory(
|
|
dry_run=dry_run, entry=entry, entry_metadata=entry_metadata
|
|
)
|
|
|
|
# Re-save the download archive after each entry is moved to the output directory
|
|
if self.maintain_download_archive:
|
|
self._enhanced_download_archive.save_download_mappings()
|
|
|
|
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):
|
|
# Break if it is None, it is indicated to not process any further
|
|
if (entry_ := plugin.modify_entry(entry_)) is None:
|
|
break
|
|
|
|
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,
|
|
)
|
|
|
|
plugins_post_split = sorted(
|
|
[plugin for plugin in plugins if plugin.priority.modify_entry_after_split],
|
|
key=lambda _plugin: _plugin.priority.modify_entry,
|
|
)
|
|
|
|
# First, modify the entry with pre_split plugins
|
|
for plugin in plugins_pre_split:
|
|
# 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
|
|
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
|
|
|
|
# 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 _process_subscription(
|
|
self,
|
|
plugins: List[Plugin],
|
|
entries: Iterable[Entry] | Iterable[Tuple[Entry, FileMetadata]],
|
|
dry_run: bool,
|
|
) -> FileHandlerTransactionLog:
|
|
for entry in entries:
|
|
entry_metadata = FileMetadata()
|
|
if isinstance(entry, tuple):
|
|
entry, entry_metadata = entry
|
|
|
|
if split_plugin := _get_split_plugin(plugins):
|
|
self._process_split_entry(
|
|
split_plugin=split_plugin, plugins=plugins, dry_run=dry_run, entry=entry
|
|
)
|
|
else:
|
|
self._process_entry(
|
|
plugins=plugins, dry_run=dry_run, entry=entry, entry_metadata=entry_metadata
|
|
)
|
|
|
|
for plugin in plugins:
|
|
plugin.post_process_subscription()
|
|
|
|
return self._enhanced_download_archive.get_file_handler_transaction_log()
|
|
|
|
def download(self, dry_run: bool = False) -> FileHandlerTransactionLog:
|
|
"""
|
|
Performs the subscription download
|
|
|
|
Parameters
|
|
----------
|
|
dry_run
|
|
If true, do not download any video/audio files or move anything to the output
|
|
directory.
|
|
"""
|
|
self._enhanced_download_archive.reinitialize(dry_run=dry_run)
|
|
plugins = self._initialize_plugins()
|
|
|
|
subscription_ytdl_options = SubscriptionYTDLOptions(
|
|
preset=self._preset_options,
|
|
plugins=plugins,
|
|
enhanced_download_archive=self._enhanced_download_archive,
|
|
working_directory=self.working_directory,
|
|
dry_run=dry_run,
|
|
)
|
|
|
|
with self._subscription_download_context_managers():
|
|
downloader = self.downloader_class(
|
|
download_options=self.downloader_options,
|
|
enhanced_download_archive=self._enhanced_download_archive,
|
|
download_ytdl_options=subscription_ytdl_options.download_builder(),
|
|
metadata_ytdl_options=subscription_ytdl_options.metadata_builder(),
|
|
overrides=self.overrides,
|
|
)
|
|
|
|
return self._process_subscription(
|
|
plugins=plugins,
|
|
entries=downloader.download(),
|
|
dry_run=dry_run,
|
|
)
|
|
|
|
def _get_entries_for_reformat(
|
|
self, download_mappings: DownloadMappings, dry_run: bool
|
|
) -> Iterable[Entry]:
|
|
entry_mapping: List[Tuple[Entry, Set[str]]] = []
|
|
for download_mapping in download_mappings._entry_mappings.values():
|
|
maybe_entry: Optional[Entry] = None
|
|
for file_name in download_mapping.file_names:
|
|
if file_name.endswith(".info.json"):
|
|
try:
|
|
with open(
|
|
Path(self.output_directory) / file_name, "r", encoding="utf-8"
|
|
) as maybe_info_json:
|
|
maybe_entry = Entry(
|
|
entry_dict=json.load(maybe_info_json),
|
|
working_directory=self.working_directory,
|
|
)
|
|
except Exception as exc:
|
|
raise ValidationException(
|
|
"info.json file cannot be loaded - subscription cannot be reformatted"
|
|
) from exc
|
|
|
|
if not maybe_entry:
|
|
raise ValidationException(
|
|
".info.json file could not be found - subscription cannot be reformatted"
|
|
)
|
|
|
|
entry_mapping.append((maybe_entry, download_mapping.file_names))
|
|
|
|
for entry, file_names in entry_mapping:
|
|
file_names_mtime: Dict[str, float] = {}
|
|
for file_name in file_names:
|
|
ext = get_file_extension(file_name)
|
|
|
|
file_path = Path(self.output_directory) / file_name
|
|
working_directory_file_path = Path(self.working_directory) / f"{entry.uid}.{ext}"
|
|
|
|
file_names_mtime[file_name] = os.path.getmtime(file_path)
|
|
|
|
# NFO files will always get rewritten, so ignore
|
|
if ext == "nfo":
|
|
continue
|
|
|
|
if not dry_run:
|
|
FileHandler.copy(
|
|
src_file_path=file_path,
|
|
dst_file_path=working_directory_file_path,
|
|
)
|
|
|
|
yield entry
|
|
|
|
for file_name, mtime in file_names_mtime.items():
|
|
# If the entry file_path is unchanged, then delete it since it was not part of the
|
|
# reformat output
|
|
if os.path.getmtime(Path(self.output_directory) / file_name) == mtime:
|
|
self._enhanced_download_archive._file_handler.delete_file_from_output_directory(
|
|
file_name
|
|
)
|
|
|
|
def update_with_info_json(self, dry_run: bool = False) -> FileHandlerTransactionLog:
|
|
plugins = self._initialize_plugins()
|
|
|
|
with self._subscription_download_context_managers():
|
|
download_mappings = self._enhanced_download_archive.mapping
|
|
self._enhanced_download_archive.reinitialize(dry_run=dry_run)
|
|
|
|
return self._process_subscription(
|
|
plugins=plugins,
|
|
entries=self._get_entries_for_reformat(
|
|
download_mappings=download_mappings, dry_run=dry_run
|
|
),
|
|
dry_run=dry_run,
|
|
)
|