9df40497e2
* dry-run huge refactor * small fixes, all tests passing * file handler, transaction log, metadata * dry-run actually dry running * working tests, still have more * Everything passing besides channel * fix youtube channel thumbnail issue, need to fix fixtures on test * lint 10, channel tests passing. Need dry-run for split and merge video * all tests passing, lint 10, should test dry-run arg
638 lines
19 KiB
Python
638 lines
19 KiB
Python
import copy
|
|
import json
|
|
import os.path
|
|
from dataclasses import dataclass
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
from typing import Any
|
|
from typing import Dict
|
|
from typing import List
|
|
from typing import Optional
|
|
from typing import Set
|
|
|
|
from yt_dlp import DateRange
|
|
|
|
from ytdl_sub.entries.entry import Entry
|
|
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.logger import Logger
|
|
|
|
|
|
@dataclass
|
|
class DownloadMapping:
|
|
upload_date: str
|
|
extractor: str
|
|
file_names: Set[str]
|
|
|
|
@property
|
|
def dict(self) -> Dict[str, Any]:
|
|
"""
|
|
:return: DownloadMapping as a dict that is serializable
|
|
"""
|
|
return {
|
|
"upload_date": self.upload_date,
|
|
"extractor": self.extractor,
|
|
"file_names": sorted(list(self.file_names)),
|
|
}
|
|
|
|
@classmethod
|
|
def from_dict(cls, mapping_dict: dict) -> "DownloadMapping":
|
|
"""
|
|
Parameters
|
|
----------
|
|
mapping_dict
|
|
Download mapping in dict format
|
|
|
|
Returns
|
|
-------
|
|
Instantiated DownloadMapping class
|
|
"""
|
|
return DownloadMapping(
|
|
upload_date=mapping_dict["upload_date"],
|
|
extractor=mapping_dict["extractor"],
|
|
file_names=set(mapping_dict["file_names"]),
|
|
)
|
|
|
|
@classmethod
|
|
def from_entry(cls, entry: Entry) -> "DownloadMapping":
|
|
"""
|
|
Parameters
|
|
----------
|
|
entry
|
|
Entry to create a download mapping for
|
|
|
|
Returns
|
|
-------
|
|
DownloadMapping for the entry
|
|
"""
|
|
return DownloadMapping(
|
|
upload_date=entry.upload_date_standardized,
|
|
extractor=entry.extractor,
|
|
file_names=set(),
|
|
)
|
|
|
|
|
|
class DownloadArchive:
|
|
"""
|
|
Class to handle any operations to the ytdl download archive. Try to keep it as barebones as
|
|
possible in case of future changes.
|
|
"""
|
|
|
|
def __init__(self, download_archive_lines: List[str]):
|
|
"""
|
|
Parameters
|
|
----------
|
|
download_archive_lines
|
|
Lines found in a YTDL download archive file, i.e. youtube id-32342343423
|
|
"""
|
|
self._download_archive_lines = download_archive_lines
|
|
|
|
@classmethod
|
|
def from_file(cls, file_path: str) -> "DownloadArchive":
|
|
"""
|
|
Parameters
|
|
----------
|
|
file_path
|
|
Path to a download archive file
|
|
|
|
Returns
|
|
-------
|
|
Instantiated DownloadArchive class
|
|
"""
|
|
# If no download archive file exists, instantiate an empty one
|
|
if not os.path.isfile(file_path):
|
|
return cls(download_archive_lines=[])
|
|
|
|
with open(file_path, "r", encoding="utf8") as file:
|
|
return cls(download_archive_lines=file.readlines())
|
|
|
|
def to_file(self, file_path: str) -> "DownloadArchive":
|
|
"""
|
|
Parameters
|
|
----------
|
|
file_path
|
|
File path to store this download archive to
|
|
|
|
Returns
|
|
-------
|
|
self
|
|
"""
|
|
with open(file_path, "w", encoding="utf8") as file:
|
|
for line in self._download_archive_lines:
|
|
file.write(f"{line}\n")
|
|
return self
|
|
|
|
def contains(self, entry_id: str) -> bool:
|
|
"""
|
|
Parameters
|
|
----------
|
|
entry_id
|
|
Id of the entry
|
|
|
|
Returns
|
|
-------
|
|
True if the entry id is within this download archive. False otherwise.
|
|
"""
|
|
return any(entry_id in line for line in self._download_archive_lines)
|
|
|
|
def remove_entry(self, entry_id: str) -> "DownloadArchive":
|
|
"""
|
|
Parameters
|
|
----------
|
|
entry_id
|
|
Entry ID to remove if it exists in this download archive
|
|
|
|
Returns
|
|
-------
|
|
self
|
|
"""
|
|
self._download_archive_lines = [
|
|
line for line in self._download_archive_lines if entry_id not in line
|
|
]
|
|
return self
|
|
|
|
|
|
class DownloadMappings:
|
|
_strptime_format = "%Y-%m-%d"
|
|
|
|
def __init__(self):
|
|
"""
|
|
Initializes an empty mapping
|
|
"""
|
|
self._entry_mappings: Dict[str, DownloadMapping] = {}
|
|
|
|
@classmethod
|
|
def from_file(cls, json_file_path: str) -> "DownloadMappings":
|
|
"""
|
|
Parameters
|
|
----------
|
|
json_file_path
|
|
Path to a json file that contains download mappings
|
|
|
|
Returns
|
|
-------
|
|
Instantiated DownloadMappings class
|
|
"""
|
|
with open(json_file_path, "r", encoding="utf8") as json_file:
|
|
entry_mappings_json = json.load(json_file)
|
|
|
|
for uid in entry_mappings_json.keys():
|
|
entry_mappings_json[uid] = DownloadMapping.from_dict(
|
|
mapping_dict=entry_mappings_json[uid]
|
|
)
|
|
|
|
download_mappings = DownloadMappings()
|
|
download_mappings._entry_mappings = entry_mappings_json
|
|
return download_mappings
|
|
|
|
@property
|
|
def entry_ids(self) -> List[str]:
|
|
"""
|
|
Returns
|
|
-------
|
|
List of entry ids in the mapping
|
|
"""
|
|
return list(self._entry_mappings.keys())
|
|
|
|
@property
|
|
def is_empty(self) -> bool:
|
|
"""
|
|
Returns
|
|
-------
|
|
True if there are no entry mappings. False otherwise.
|
|
"""
|
|
return len(self._entry_mappings) == 0
|
|
|
|
def add_entry(self, entry: Entry, entry_file_path: str) -> "DownloadMappings":
|
|
"""
|
|
Adds a file path for the entry. An entry can map to multiple file paths.
|
|
|
|
Parameters
|
|
----------
|
|
entry
|
|
Entry that this file belongs to
|
|
entry_file_path
|
|
Relative path to the file that lives in the output directory
|
|
|
|
Returns
|
|
-------
|
|
self
|
|
"""
|
|
if entry.uid not in self.entry_ids:
|
|
self._entry_mappings[entry.uid] = DownloadMapping.from_entry(entry=entry)
|
|
|
|
self._entry_mappings[entry.uid].file_names.add(entry_file_path)
|
|
return self
|
|
|
|
def remove_entry(self, entry_id: str) -> "DownloadMappings":
|
|
"""
|
|
Parameters
|
|
----------
|
|
entry_id
|
|
Id of the entry to remove
|
|
|
|
Returns
|
|
-------
|
|
self
|
|
"""
|
|
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]:
|
|
"""
|
|
Parameters
|
|
----------
|
|
date_range
|
|
Range of dates that entries' upload dates must be within
|
|
|
|
Returns
|
|
-------
|
|
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(
|
|
out_of_range_entry_mappings[uid].upload_date,
|
|
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":
|
|
"""
|
|
Parameters
|
|
----------
|
|
output_json_file
|
|
Output json file path to write the download mappings to
|
|
|
|
Returns
|
|
-------
|
|
self
|
|
"""
|
|
# Create json string first to ensure it is valid before writing anything to file
|
|
json_str = json.dumps(
|
|
obj={
|
|
uid: mapping.dict
|
|
for uid, mapping in sorted(
|
|
self._entry_mappings.items(),
|
|
key=lambda item: item[1].upload_date,
|
|
reverse=True,
|
|
)
|
|
},
|
|
indent=2,
|
|
sort_keys=True,
|
|
)
|
|
|
|
with open(output_json_file, "w", encoding="utf8") as file:
|
|
file.write(json_str)
|
|
|
|
return self
|
|
|
|
def to_download_archive(self) -> DownloadArchive:
|
|
"""
|
|
Returns
|
|
-------
|
|
A DownloadArchive created from the DownloadMappings' ids and extractors. YTDL will use this
|
|
to avoid redownloading entries already downloaded.
|
|
"""
|
|
lines: List[str] = []
|
|
for entry_id, metadata in self._entry_mappings.items():
|
|
lines.append(f"{metadata.extractor} {entry_id}")
|
|
|
|
return DownloadArchive(download_archive_lines=lines)
|
|
|
|
|
|
class EnhancedDownloadArchive:
|
|
"""
|
|
Maintains ytdl's download archive file as well as create an additional mapping file to map
|
|
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 an existing enhanced download archive file
|
|
exists. If so, load it into the class. Otherwise, initialize an empty instance of one.
|
|
b. self._copy_to_working_directory()
|
|
If the download archive was loaded successfully, create a ytdl download archive in the
|
|
working directory. This will let ytdl know which files are already downloaded.
|
|
2. ( Perform the ytdlp download using a download archive with the same name )
|
|
a. An existing archive should have been copied into the working directory for reuse
|
|
3. self.mapping.add_entry(entry, file_path)
|
|
a. Should be called for any file created for the given entry that gets moved to the output
|
|
directory
|
|
4. OPTIONAL: self.remove_stale_files()
|
|
a. After all files have been moved over in the output directory, remove any stale files that
|
|
exist in there.
|
|
5. self.save_download_archive()
|
|
a. Save the updated mapping file to the output directory.
|
|
6. ( Delete the working directory )
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
subscription_name: str,
|
|
working_directory: str,
|
|
output_directory: str,
|
|
dry_run: bool = False,
|
|
):
|
|
self.subscription_name = subscription_name
|
|
self._file_handler = FileHandler(
|
|
working_directory=working_directory, output_directory=output_directory, dry_run=dry_run
|
|
)
|
|
self._download_mapping = DownloadMappings()
|
|
|
|
self._logger = Logger.get(name=subscription_name)
|
|
|
|
def reinitialize(self, dry_run: bool) -> "EnhancedDownloadArchive":
|
|
"""
|
|
Re-initialize the enhanced download archive for successive downloads w/the same
|
|
subscription.
|
|
|
|
Parameters
|
|
----------
|
|
dry_run
|
|
Whether to actually move files to the output directory
|
|
|
|
Returns
|
|
-------
|
|
self
|
|
"""
|
|
self._file_handler = FileHandler(
|
|
working_directory=self.working_directory,
|
|
output_directory=self.output_directory,
|
|
dry_run=dry_run,
|
|
)
|
|
self._download_mapping = DownloadMappings()
|
|
return self
|
|
|
|
@property
|
|
def is_dry_run(self) -> bool:
|
|
"""
|
|
Returns
|
|
-------
|
|
True if this session is a dry-run. False otherwise.
|
|
"""
|
|
return self._file_handler.dry_run
|
|
|
|
@property
|
|
def archive_file_name(self) -> str:
|
|
"""
|
|
Returns
|
|
-------
|
|
The download archive's file name (no path)
|
|
"""
|
|
return f".ytdl-subscribe-{self.subscription_name}-download-archive.txt"
|
|
|
|
@property
|
|
def working_directory(self) -> str:
|
|
"""
|
|
Returns
|
|
-------
|
|
Path to the working directory
|
|
"""
|
|
return self._file_handler.working_directory
|
|
|
|
@property
|
|
def output_directory(self) -> str:
|
|
"""
|
|
Returns
|
|
-------
|
|
Path to the output directory
|
|
"""
|
|
return self._file_handler.output_directory
|
|
|
|
@property
|
|
def _mapping_file_name(self) -> str:
|
|
"""
|
|
Returns
|
|
-------
|
|
The download mapping's file name (no path)
|
|
"""
|
|
return f".ytdl-sub-{self.subscription_name}-download-archive.json"
|
|
|
|
@property
|
|
def _mapping_output_file_path(self) -> str:
|
|
"""
|
|
Returns
|
|
-------
|
|
The download mapping's file path in the output directory.
|
|
"""
|
|
return str(Path(self.output_directory) / self._mapping_file_name)
|
|
|
|
@property
|
|
def _mapping_working_file_path(self) -> str:
|
|
"""
|
|
Returns
|
|
-------
|
|
The download mapping's file path in the working directory.
|
|
"""
|
|
return str(Path(self.working_directory) / self._mapping_file_name)
|
|
|
|
@property
|
|
def _archive_working_file_path(self) -> str:
|
|
"""
|
|
Returns
|
|
-------
|
|
The download archive's file path in the working directory.
|
|
"""
|
|
return str(Path(self.working_directory) / self.archive_file_name)
|
|
|
|
@property
|
|
def mapping(self) -> DownloadMappings:
|
|
"""
|
|
Returns
|
|
-------
|
|
Loaded DownloadMappings
|
|
|
|
Raises
|
|
------
|
|
ValueError
|
|
If the download mappings was not loaded
|
|
"""
|
|
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":
|
|
"""
|
|
Tries to load download mappings if they are present in the output directory.
|
|
|
|
Returns
|
|
-------
|
|
self
|
|
"""
|
|
# If a mapping file exists in the output directory, load it up.
|
|
if os.path.isfile(self._mapping_output_file_path):
|
|
self._download_mapping = DownloadMappings.from_file(
|
|
json_file_path=self._mapping_output_file_path
|
|
)
|
|
return self
|
|
|
|
def _copy_mapping_to_working_directory(self) -> "EnhancedDownloadArchive":
|
|
"""
|
|
If the mapping is not empty, create a download archive from it and save it into the
|
|
working directory. This will tell YTDL to not redownload already downloaded entries.
|
|
|
|
Returns
|
|
-------
|
|
self
|
|
"""
|
|
# If the download mapping is empty, do nothing since the ytdl downloader will create a new
|
|
# download archive file
|
|
if self.mapping.is_empty:
|
|
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)
|
|
|
|
return self
|
|
|
|
def prepare_download_archive(self) -> "EnhancedDownloadArchive":
|
|
"""
|
|
Helper function to load mappings and create a YTDL download archive file in the
|
|
working directory if mappings exist.
|
|
|
|
Returns
|
|
-------
|
|
self
|
|
"""
|
|
self._load()
|
|
self._copy_mapping_to_working_directory()
|
|
return self
|
|
|
|
def remove_stale_files(self, date_range: DateRange) -> "EnhancedDownloadArchive":
|
|
"""
|
|
Checks all entries within the mappings. If any entries' upload dates are not within the
|
|
provided date range, delete them.
|
|
|
|
Parameters
|
|
----------
|
|
date_range
|
|
Date range the upload date must be in to not get deleted
|
|
|
|
Returns
|
|
-------
|
|
self
|
|
"""
|
|
stale_mappings: Dict[str, DownloadMapping] = self.mapping.get_entries_out_of_range(
|
|
date_range=date_range
|
|
)
|
|
|
|
for uid, mapping in stale_mappings.items():
|
|
for file_name in mapping.file_names:
|
|
self._file_handler.delete_file_from_output_directory(file_name=file_name)
|
|
|
|
self.mapping.remove_entry(entry_id=uid)
|
|
|
|
return self
|
|
|
|
def save_download_mappings(self) -> "EnhancedDownloadArchive":
|
|
"""
|
|
Saves the updated download mappings to the output directory
|
|
|
|
Returns
|
|
-------
|
|
self
|
|
"""
|
|
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
|
|
|
|
def save_file_to_output_directory(
|
|
self,
|
|
file_name: str,
|
|
file_metadata: Optional[FileMetadata] = None,
|
|
output_file_name: Optional[str] = None,
|
|
entry: Optional[Entry] = None,
|
|
):
|
|
"""
|
|
Saves a file from the working directory to the output directory and record it in the
|
|
transaction log.
|
|
|
|
Parameters
|
|
----------
|
|
file_name
|
|
Name of the file to move (does not include working directory path)
|
|
file_metadata
|
|
Optional. Metadata to record to the transaction log for this file
|
|
output_file_name
|
|
Optional. Final name of the file in the output directory (does not include output
|
|
directory path). If None, use the same working_directory file_name
|
|
entry
|
|
Optional. Entry that this file belongs to
|
|
"""
|
|
if output_file_name is None:
|
|
output_file_name = file_name
|
|
|
|
if entry:
|
|
self.mapping.add_entry(entry=entry, entry_file_path=output_file_name)
|
|
|
|
self._file_handler.copy_file_to_output_directory(
|
|
file_name=file_name, file_metadata=file_metadata, output_file_name=output_file_name
|
|
)
|
|
|
|
def get_file_handler_transaction_log(self) -> FileHandlerTransactionLog:
|
|
"""
|
|
Returns
|
|
-------
|
|
File handler transaction log for this session
|
|
"""
|
|
return self._file_handler.file_handler_transaction_log
|
|
|
|
|
|
class DownloadArchiver:
|
|
"""
|
|
Used for any class that saves files. Does not allow direct access to output_directory,
|
|
forcing the user of the class to use ``save_file`` so it gets archived and avoids any writes
|
|
during dry-run.
|
|
"""
|
|
|
|
def __init__(self, enhanced_download_archive: EnhancedDownloadArchive):
|
|
self.__enhanced_download_archive = enhanced_download_archive
|
|
|
|
@property
|
|
def working_directory(self) -> str:
|
|
"""
|
|
Returns
|
|
-------
|
|
Path to the working directory
|
|
"""
|
|
return self.__enhanced_download_archive.working_directory
|
|
|
|
@property
|
|
def is_dry_run(self) -> bool:
|
|
"""
|
|
Returns
|
|
-------
|
|
True if this session is a dry-run. False otherwise.
|
|
"""
|
|
return self.__enhanced_download_archive.is_dry_run
|
|
|
|
def save_file(
|
|
self, file_name: str, output_file_name: Optional[str] = None, entry: Optional[Entry] = None
|
|
) -> None:
|
|
"""
|
|
Saves a file in the working directory to the output directory.
|
|
|
|
Parameters
|
|
----------
|
|
file_name
|
|
Name of the file relative to the working directory
|
|
output_file_name
|
|
Optional. Final name of the file in the output directory (does not include output
|
|
directory path). If None, use the same working_directory file_name
|
|
entry
|
|
Optional. Entry that the file belongs to
|
|
"""
|
|
self.__enhanced_download_archive.save_file_to_output_directory(
|
|
file_name=file_name, output_file_name=output_file_name, entry=entry
|
|
)
|