ytdl-sub/src/ytdl_sub/cli/entrypoint.py
Jesse Bannon 07e5cd008c all ruff
2026-03-09 11:22:00 -07:00

288 lines
9.4 KiB
Python

import gc
import os
import random
import sys
from datetime import datetime
from pathlib import Path
from typing import Dict, List, Optional
from yt_dlp.utils import sanitize_filename
from ytdl_sub.cli.output_summary import output_summary
from ytdl_sub.cli.output_transaction_log import (
_maybe_validate_transaction_log_file,
output_transaction_log,
)
from ytdl_sub.cli.parsers.cli_to_sub import print_cli_to_sub
from ytdl_sub.cli.parsers.dl import DownloadArgsParser
from ytdl_sub.cli.parsers.main import DEFAULT_CONFIG_FILE_NAME, parser
from ytdl_sub.config.config_file import ConfigFile
from ytdl_sub.subscriptions.subscription import Subscription
from ytdl_sub.utils.exceptions import ExperimentalFeatureNotEnabled, ValidationException
from ytdl_sub.utils.file_handler import FileHandler
from ytdl_sub.utils.file_lock import working_directory_lock
from ytdl_sub.utils.logger import Logger
logger = Logger.get()
# View is a command to run a simple dry-run on a URL using the `_view` preset.
# Use ytdl-sub dl arguments to use the preset
_VIEW_EXTRA_ARGS_FORMATTER = "--preset _view --overrides.url {}"
def _log_time() -> str:
return datetime.now().strftime("%Y-%m-%d-%H%M%S")
def _maybe_write_subscription_log_file(
config: ConfigFile,
subscription: Subscription,
dry_run: bool,
exception: Optional[Exception] = None,
) -> None:
success: bool = exception is None
# If dry-run, do nothing
if dry_run:
return
# If persisting logs is disabled, do nothing
if not config.config_options.persist_logs:
return
# If persisting successful logs is disabled, do nothing
if success and not config.config_options.persist_logs.keep_successful_logs:
return
log_subscription_name = sanitize_filename(subscription.name).lower().replace(" ", "_")
log_success = "success" if success else "error"
log_filename = f"{_log_time()}.{log_subscription_name}.{log_success}.log"
persist_log_path = Path(config.config_options.persist_logs.logs_directory) / log_filename
if not success:
Logger.log_exception(exception=exception, log_filepath=persist_log_path)
os.makedirs(os.path.dirname(persist_log_path), exist_ok=True)
FileHandler.copy(Logger.debug_log_filename(), persist_log_path)
def _download_subscriptions_from_yaml_files(
config: ConfigFile,
subscription_paths: List[str],
subscription_matches: List[str],
subscription_override_dict: Dict,
update_with_info_json: bool,
dry_run: bool,
shuffle: bool,
) -> List[Subscription]:
"""
Downloads all subscriptions from one or many subscription yaml files.
Parameters
----------
config
Configuration file
subscription_paths
Path to subscription files to download
subscription_matches
Optional list of substrings to match subscription names to (only run if matched)
update_with_info_json
Whether to actually download or update using existing info json
dry_run
Whether to dry run or not
shuffle
Whether to shuffle the subscription download order
Returns
-------
List of subscriptions processed
Raises
------
Exception
Any exception during download
"""
subscriptions: List[Subscription] = []
# Load all the subscriptions first to perform all validation before downloading
for path in subscription_paths:
subscriptions += Subscription.from_file_path(
config=config,
subscription_path=path,
subscription_matches=subscription_matches,
subscription_override_dict=subscription_override_dict,
)
if shuffle:
logger.info("Shuffling subscriptions")
random.shuffle(subscriptions)
for subscription in subscriptions:
with subscription.exception_handling():
logger.info(
"Beginning subscription %s for %s",
("dry run" if dry_run else "download"),
subscription.name,
)
logger.debug("Subscription full yaml:\n%s", subscription.as_yaml())
if update_with_info_json:
subscription.update_with_info_json(dry_run=dry_run)
else:
subscription.download(dry_run=dry_run)
_maybe_write_subscription_log_file(
config=config,
subscription=subscription,
dry_run=dry_run,
exception=subscription.exception,
)
Logger.cleanup(has_error=False)
gc.collect() # Garbage collect after each subscription download
return subscriptions
def _download_subscription_from_cli(
config: ConfigFile, dry_run: bool, extra_args: List[str]
) -> Subscription:
"""
Downloads a one-off subscription using the CLI
Parameters
----------
config
Configuration file
dry_run
Whether this is a dry-run
extra_args
Extra arguments from argparse that contain dynamic subscription options
Returns
-------
Subscription and its download transaction log
"""
dl_args_parser = DownloadArgsParser(
extra_arguments=extra_args, config_options=config.config_options
)
subscription_args_dict = dl_args_parser.to_subscription_dict()
subscription_name = dl_args_parser.get_dl_subscription_name()
subscription = Subscription.from_dict(
config=config, preset_name=subscription_name, preset_dict=subscription_args_dict
)
logger.info("Beginning CLI %s", ("dry run" if dry_run else "download"))
subscription.download(dry_run=dry_run)
return subscription
def _view_url_from_cli(config: ConfigFile, url: str, split_chapters: bool) -> Subscription:
"""
`ytdl-sub view` dry-runs a URL to print its source variables. Use the pre-built `_view` preset,
inject the URL argument, and dry-run.
"""
preset = "_view_split_chapters" if split_chapters else "_view"
subscription = Subscription.from_dict(
config=config,
preset_name="ytdl-sub-view",
preset_dict={"preset": preset, "overrides": {"url": url}},
)
logger.info(
"Viewing source variables for URL '%s'%s",
url,
" with split chapters" if split_chapters else "",
)
subscription.download(dry_run=True)
return subscription
def main() -> List[Subscription]:
"""
Entrypoint for ytdl-sub, without the error handling
"""
# If no args are provided, print help and exit
if len(sys.argv) < 2:
parser.print_help()
return []
args, extra_args = parser.parse_known_args()
if args.subparser == "cli-to-sub":
print_cli_to_sub(args=extra_args)
return []
# Load the config
if args.config:
config = ConfigFile.from_file_path(args.config)
elif os.path.isfile(DEFAULT_CONFIG_FILE_NAME):
config = ConfigFile.from_file_path(DEFAULT_CONFIG_FILE_NAME)
else:
logger.info("No config specified, using defaults.")
config = ConfigFile.default()
subscriptions: List[Subscription] = []
# If transaction log file is specified, make sure we can open it
_maybe_validate_transaction_log_file(transaction_log_file_path=args.transaction_log)
with working_directory_lock(config=config):
if args.subparser == "sub":
if (
args.update_with_info_json
and not config.config_options.experimental.enable_update_with_info_json
):
raise ExperimentalFeatureNotEnabled(
"--update-with-info-json requires setting"
" configuration.experimental.enable_update_with_info_json to True. This"
" feature is ",
"still being tested and has the ability to destroy files. Ensure you have a ",
"full backup before usage. You have been warned!",
)
subscription_override_dict = {}
if args.dl_override:
subscription_override_dict = DownloadArgsParser.from_dl_override(
override=args.dl_override, config=config
).to_subscription_dict()
logger.info("Validating subscriptions...")
subscriptions = _download_subscriptions_from_yaml_files(
config=config,
subscription_paths=args.subscription_paths,
subscription_matches=args.match,
subscription_override_dict=subscription_override_dict,
update_with_info_json=args.update_with_info_json,
dry_run=args.dry_run,
shuffle=args.shuffle,
)
# One-off download
elif args.subparser == "dl":
logger.info("Validating presets...")
subscriptions.append(
_download_subscription_from_cli(
config=config, dry_run=args.dry_run, extra_args=extra_args
)
)
elif args.subparser == "view":
subscriptions.append(
_view_url_from_cli(config=config, url=args.url, split_chapters=args.split_chapters)
)
else:
raise ValidationException("Must provide one of the commands: sub, dl, view, cli-to-sub")
if not args.suppress_transaction_log:
output_transaction_log(
subscriptions=subscriptions,
transaction_log_file_path=args.transaction_log,
)
output_summary(subscriptions, suppress_colors=args.suppress_colors)
return subscriptions