diff --git a/ytdl_subscribe/validators/exceptions.py b/ytdl_subscribe/validators/exceptions.py index f787481d..0f20ed14 100644 --- a/ytdl_subscribe/validators/exceptions.py +++ b/ytdl_subscribe/validators/exceptions.py @@ -8,3 +8,7 @@ class StringFormattingException(ValidationException): class StringFormattingVariableNotFoundException(StringFormattingException): """Tried to format a string but the variable was not found""" + + +class DownloadArchiveException(ValueError): + """Any user or file errors caused by download archive or mapping files""" diff --git a/ytdl_subscribe/ytdl_additions/enhanced_download_archive.py b/ytdl_subscribe/ytdl_additions/enhanced_download_archive.py index ad45ad71..d81e2c9d 100644 --- a/ytdl_subscribe/ytdl_additions/enhanced_download_archive.py +++ b/ytdl_subscribe/ytdl_additions/enhanced_download_archive.py @@ -1,39 +1,95 @@ +import copy import json +import os.path +import shutil from dataclasses import dataclass +from datetime import datetime from pathlib import Path from typing import Dict from typing import List +from typing import Optional +from typing import Set + +from yt_dlp import DateRange from ytdl_subscribe.entries.entry import Entry +from ytdl_subscribe.validators.exceptions import DownloadArchiveException @dataclass class DownloadMapping: upload_date: str - file_paths: List[str] + file_paths: Set[str] class DownloadMappings: + _strptime_format = "%Y-%m-%d" + def __init__(self): self._entry_mappings: Dict[str, DownloadMapping] = {} @classmethod - def from_json(cls, file_path: str) -> "DownloadMappings": - entry_mappings = DownloadMappings() - entry_mappings._entry_mappings = json.load(open(file_path, "r", encoding="utf8")) - return entry_mappings + def from_file(cls, json_file_path: str) -> "DownloadMappings": + entry_mappings_json = json.load(open(json_file_path, "r", encoding="utf8")) + for uid in entry_mappings_json.keys(): + entry_mappings_json[uid] = DownloadMapping( + upload_date=entry_mappings_json[uid]["upload_date"], + file_paths=set(entry_mappings_json[uid]["file_paths"]), + ) - def to_json(self) -> str: - return json.dumps( - { + download_mappings = DownloadMappings() + download_mappings._entry_mappings = entry_mappings_json + return download_mappings + + @property + def entry_ids(self) -> List[str]: + return list(self._entry_mappings.keys()) + + def add_entry(self, entry: Entry, entry_file_path: str) -> "DownloadMappings": + if entry.uid in self.entry_ids: + self._entry_mappings[entry.uid].file_paths.add(entry_file_path) + else: + self._entry_mappings[entry.uid] = DownloadMapping( + upload_date=entry.standardized_upload_date, file_paths={entry_file_path} + ) + return self + + def remove_entry(self, entry_id: str) -> "DownloadMappings": + if entry_id in self.entry_ids: + del self._entry_mappings[entry_id] + return self + + def get_entries_out_of_range(self, date_range: DateRange) -> Dict[str, DownloadMapping]: + """ + :param date_range: range of dates that entries' upload dates must be within + :return: dict of entry_id: mapping if the upload date is not in the date range + """ + out_of_range_entry_mappings = copy.deepcopy(self._entry_mappings) + for uid in list(out_of_range_entry_mappings.keys()): + upload_date = datetime.strptime( + date_string=out_of_range_entry_mappings[uid].upload_date, + format=self._strptime_format, + ).date() + + if upload_date in date_range: + del out_of_range_entry_mappings[uid] + + return out_of_range_entry_mappings + + def to_file(self, output_json_file: str) -> "DownloadMappings": + json.dump( + obj={ uid: mapping for uid, mapping in sorted( self._entry_mappings.items(), key=lambda item: item[1]["upload_date"], reverse=True, ) - } + }, + fp=open(output_json_file, "w", encoding="utf8"), + indent=2, ) + return self class DownloadArchive: @@ -42,9 +98,7 @@ class DownloadArchive: possible in case of future changes. """ - def __init__( - self, - ): + def __init__(self): self._download_archive_lines: List[str] = [] @classmethod @@ -54,39 +108,179 @@ class DownloadArchive: download_archive._download_archive_lines = lines return download_archive + def to_file(self, file_path: str) -> "DownloadArchive": + with open(file_path, "w", encoding="utf8") as file: + file.writelines(self._download_archive_lines) + return self + + def contains(self, entry_id: str) -> bool: + return any(entry_id in line for line in self._download_archive_lines) + + def remove_entry(self, entry_id: str) -> "DownloadArchive": + self._download_archive_lines = [ + line for line in self._download_archive_lines if entry_id not in line + ] + return self + class EnhancedDownloadArchive: """ Maintains ytdl's download archive file as well as create an additional mapping file to map - ytdl ids to multiple files + ytdl ids to multiple files. Used to delete 'stale' files that are out of range based on the + file's entry's upload date. + + Should be used in this order: + + 1. self.prepare_download_archive() + Internally calls... + a. self._load() + - Checks the output directory to see if existing download archive/mapping files exist. + Loads them into the class if they do. If both do not exist, do not load anything since + they are getting instantiated for the first time. If only one archive/mapping file + exists but not the other, error. + b. self._copy_to_working_directory() + - If the files were loaded successfully, copy them to the working directory. This will + let ytdlp reuse the download archive file. + 2. ( Perform the ytdlp download using a download archive with the same name ) + - An existing archive should have been copied into the working directory for reuse + 3. self.mapping.add_entry(entry, file_path) + - Should be called for any file created for the given entry that gets moved to the output + directory + 4. OPTIONAL: self.remove_stale_files() + - After all files have been moved over in the output directory, remove any stale files that + exist in there. + 5. self.save_download_archive() + - Save the updated archive/mapping files to the output directory. + 6. ( Delete the working directory ) """ - def __init__(self, name: str): - self.name = name - self._download_archive = None - self._download_mapping = None - self._appended_entries: List[Entry] = [] + def __init__(self, subscription_name: str, working_directory: str, output_directory: str): + self.subscription_name = subscription_name + self.working_directory = working_directory + self.output_directory = output_directory + + self._download_archive: Optional[DownloadArchive] = None + self._download_mapping: Optional[DownloadMappings] = None @property - def download_archive_file_name(self): - return f".ytdl-subscribe-{self.name}-download-archive.txt" + def _archive_file_name(self) -> str: + """ + :return: The download archive's file name (no path) + """ + return f".ytdl-subscribe-{self.subscription_name}-download-archive.txt" @property - def download_mapping_file_name(self): - return f".ytdl-subscribe-{self.name}-download-mapping.json" + def _mapping_file_name(self) -> str: + """ + :return: The download mapping's file name (no path) + """ + return f".ytdl-subscribe-{self.subscription_name}-download-mapping.json" - def load(self, directory: str) -> "EnhancedDownloadArchive": - self._download_archive = DownloadArchive.from_file( - file_path=str(Path(directory) / self.download_archive_file_name) - ) + @property + def _archive_output_file_path(self) -> str: + """ + :return: The download archive's file path in the output directory. + """ + return str(Path(self.output_directory) / self._archive_file_name) - self._download_mapping = DownloadMappings.from_json( - file_path=str(Path(directory) / self.download_mapping_file_name) - ) + @property + def _mapping_output_file_path(self): + """ + :return: The download mapping's file path in the output directory. + """ + return str(Path(self.output_directory) / self._mapping_file_name) + + @property + def _archive_working_file_path(self) -> str: + """ + :return: The download archive's file path in the working directory. + """ + return str(Path(self.working_directory) / self._archive_file_name) + + @property + def _mapping_working_file_path(self): + """ + :return: The download mapping's file path in the working directory. + """ + return str(Path(self.working_directory) / self._mapping_file_name) + + @property + def mapping(self) -> DownloadMappings: + if self._download_mapping is None: + raise ValueError("Tried to use download mapping before it was loaded") + return self._download_mapping + + def load(self) -> "EnhancedDownloadArchive": + download_archive_is_file = os.path.isfile(self._mapping_output_file_path) + download_mapping_is_file = os.path.isfile(self._archive_output_file_path) + + if download_archive_is_file ^ download_mapping_is_file: + found_file_path, missing_name, missing_file_path = ( + (self._archive_output_file_path, "mapping", self._mapping_output_file_path) + if download_archive_is_file + else (self._mapping_output_file_path, "archive", self._archive_output_file_path) + ) + + raise DownloadArchiveException( + f"Error when loading archive files for {self.subscription_name}. Partially loaded " + f"the enhanced download archive, but is missing the {missing_name} file located at " + f"{missing_file_path}. Either disable 'maintain_download_archive' or delete " + f"{found_file_path} to recreate the two missing files." + ) + + # If both files exist, load them up + if download_archive_is_file and download_mapping_is_file: + self._download_archive = DownloadArchive.from_file( + file_path=self._archive_output_file_path + ) + self._download_mapping = DownloadMappings.from_file( + json_file_path=self._mapping_output_file_path + ) + # Otherwise, init an empty download mappings object. Keep _download_archive as None to + # indicate it was not loaded + else: + self._download_mapping = DownloadMappings() return self - # def append(self, entry: Entry, relative_file_path: str): - # self._appended_entries.append( - # EntryMapping(entry=entry, relative_file_path=relative_file_path)) - # + def copy_to_working_directory(self) -> "EnhancedDownloadArchive": + # If no download archive was found, then there is nothing to copy + if not self._download_archive: + return self + + shutil.copyfile(self._archive_output_file_path, self._archive_working_file_path) + shutil.copyfile(self._mapping_output_file_path, self._mapping_working_file_path) + + return self + + def remove_stale_files(self, date_range: DateRange) -> "EnhancedDownloadArchive": + stale_mappings: Dict[str, DownloadMapping] = self.mapping.get_entries_out_of_range( + date_range=date_range + ) + + for uid, mapping in stale_mappings.items(): + print(f"[{uid}] Removing the following stale file(s):") + for file_path in mapping.file_paths: + print(f" - {file_path}") + if os.path.exists(file_path): + os.remove(file_path) + + self.mapping.remove_entry(entry_id=uid) + + return self + + def save_download_archive(self) -> "EnhancedDownloadArchive": + # Load the download archive from the working directory, which should contain any past + # and new entries downloaded in this session + self._download_archive = DownloadArchive.from_file(self._archive_working_file_path) + + # Keep the download archive in sync with the mapping + for entry_id in self.mapping.entry_ids: + if not self._download_archive.contains(entry_id): + self._download_archive.remove_entry(entry_id) + + # Save the archive/mappign files + self._download_archive.to_file(file_path=self._archive_output_file_path) + self._download_mapping.to_file(output_json_file=self._mapping_output_file_path) + + return self