big upgrade UwU
This commit is contained in:
@@ -0,0 +1,349 @@
|
||||
"""Command-line interface."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import os
|
||||
import sys
|
||||
from collections.abc import Callable, Sequence
|
||||
from pathlib import Path
|
||||
|
||||
import aiohttp
|
||||
from dotenv import load_dotenv
|
||||
from telegram import Bot
|
||||
from telegram.error import InvalidToken, TelegramError
|
||||
|
||||
from sticker_downloader import __version__
|
||||
from sticker_downloader.config import (
|
||||
ALL_FILE_TYPES,
|
||||
DEFAULT_CONCURRENCY,
|
||||
DEFAULT_OUTPUT_DIR,
|
||||
DEFAULT_RETRIES,
|
||||
DEFAULT_URL_FILE,
|
||||
FILE_TYPE_HELP,
|
||||
DownloadConfig,
|
||||
parse_file_types,
|
||||
)
|
||||
from sticker_downloader.downloader import StickerDownloader
|
||||
from sticker_downloader.errors import StickerDownloaderError
|
||||
from sticker_downloader.fetcher import AiohttpFetcher
|
||||
from sticker_downloader.progress import ConsoleProgress, NullProgress, ProgressReporter
|
||||
from sticker_downloader.results import PackResult
|
||||
from sticker_downloader.urls import parse_pack_reference_lines, read_pack_references
|
||||
|
||||
EXIT_OK = 0
|
||||
EXIT_FAILURES = 1
|
||||
EXIT_USAGE = 2
|
||||
|
||||
TOKEN_ENV_VAR = "TELEGRAM_BOT_TOKEN"
|
||||
|
||||
_TOKEN_INSTRUCTIONS = f"""\
|
||||
Create a bot with @BotFather on Telegram, then either:
|
||||
* put {TOKEN_ENV_VAR}=<your token> in a .env file next to this project, or
|
||||
* export {TOKEN_ENV_VAR}=<your token>, or
|
||||
* pass --token <your token>."""
|
||||
|
||||
_TOKEN_MISSING = f"No Telegram bot token found.\n\n{_TOKEN_INSTRUCTIONS}"
|
||||
_TOKEN_REJECTED = f"Telegram rejected the bot token.\n\n{_TOKEN_INSTRUCTIONS}"
|
||||
|
||||
|
||||
class UsageError(StickerDownloaderError):
|
||||
"""Raised for bad invocations; reported without a traceback."""
|
||||
|
||||
|
||||
def build_parser() -> argparse.ArgumentParser:
|
||||
types_list = ", ".join(f"{name} ({FILE_TYPE_HELP[name]})" for name in ALL_FILE_TYPES)
|
||||
parser = argparse.ArgumentParser(
|
||||
prog="tg-stickers",
|
||||
description="Download Telegram sticker packs into per-pack folders.",
|
||||
epilog=(
|
||||
"examples:\n"
|
||||
" tg-stickers https://t.me/addstickers/SomePack\n"
|
||||
" tg-stickers SomePack AnotherPack --types webp,png\n"
|
||||
f" tg-stickers --from-file {DEFAULT_URL_FILE} --concurrency 16\n"
|
||||
" tg-stickers --from-file - < my-packs.txt\n"
|
||||
" tg-stickers # interactive prompts\n"
|
||||
),
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
)
|
||||
parser.add_argument(
|
||||
"packs",
|
||||
nargs="*",
|
||||
metavar="PACK",
|
||||
help="sticker pack URLs, tg:// links or bare pack names",
|
||||
)
|
||||
parser.add_argument(
|
||||
"-f",
|
||||
"--from-file",
|
||||
metavar="PATH",
|
||||
help=f"read references from a file, one per line ('-' for stdin); "
|
||||
f"blank lines and # comments are ignored (default: {DEFAULT_URL_FILE} "
|
||||
f"when chosen interactively)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"-t",
|
||||
"--types",
|
||||
metavar="LIST",
|
||||
help=f"comma-separated file types to save, or 'all'. Supported: {types_list}",
|
||||
)
|
||||
parser.add_argument(
|
||||
"-o",
|
||||
"--output",
|
||||
metavar="DIR",
|
||||
type=Path,
|
||||
default=DEFAULT_OUTPUT_DIR,
|
||||
help=f"directory to write packs into (default: {DEFAULT_OUTPUT_DIR})",
|
||||
)
|
||||
parser.add_argument(
|
||||
"-c",
|
||||
"--concurrency",
|
||||
type=int,
|
||||
default=DEFAULT_CONCURRENCY,
|
||||
metavar="N",
|
||||
help=f"files to download in parallel per pack (default: {DEFAULT_CONCURRENCY})",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--retries",
|
||||
type=int,
|
||||
default=DEFAULT_RETRIES,
|
||||
metavar="N",
|
||||
help=f"attempts per network call (default: {DEFAULT_RETRIES})",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--overwrite",
|
||||
action="store_true",
|
||||
help="re-download files that are already on disk (default: skip them)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--emoji-names",
|
||||
action="store_true",
|
||||
help="include each sticker's emoji in its file name",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--no-metadata",
|
||||
dest="write_metadata",
|
||||
action="store_false",
|
||||
help="do not write pack.json alongside the stickers",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--no-pack-thumbnail",
|
||||
dest="pack_thumbnail",
|
||||
action="store_false",
|
||||
help="do not download the pack's own cover thumbnail",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--dry-run",
|
||||
action="store_true",
|
||||
help="list what would be downloaded without writing anything",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--token",
|
||||
metavar="TOKEN",
|
||||
help=f"bot token (default: ${TOKEN_ENV_VAR}, also read from --env-file)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--env-file",
|
||||
metavar="PATH",
|
||||
default=".env",
|
||||
help="dotenv file to load the token from (default: .env)",
|
||||
)
|
||||
verbosity = parser.add_mutually_exclusive_group()
|
||||
verbosity.add_argument(
|
||||
"-v", "--verbose", action="store_true", help="log every file individually"
|
||||
)
|
||||
verbosity.add_argument(
|
||||
"-q", "--quiet", action="store_true", help="only report fatal errors"
|
||||
)
|
||||
parser.add_argument("--version", action="version", version=f"%(prog)s {__version__}")
|
||||
return parser
|
||||
|
||||
|
||||
def _resolve_references(
|
||||
args: argparse.Namespace, progress: ProgressReporter
|
||||
) -> list[str]:
|
||||
"""Work out which packs to download, prompting only when nothing was given."""
|
||||
if args.packs and args.from_file:
|
||||
raise UsageError("pass either PACK arguments or --from-file, not both")
|
||||
|
||||
if args.packs:
|
||||
references, problems = parse_pack_reference_lines(args.packs, source="argument")
|
||||
for problem in problems:
|
||||
progress.warn(problem.split(": ", 1)[-1])
|
||||
if not references:
|
||||
raise UsageError("no valid sticker pack references given")
|
||||
return references
|
||||
|
||||
if args.from_file:
|
||||
return _references_from_file(args.from_file, progress)
|
||||
|
||||
return []
|
||||
|
||||
|
||||
def _references_from_file(source: str, progress: ProgressReporter) -> list[str]:
|
||||
if source == "-":
|
||||
references, problems = parse_pack_reference_lines(
|
||||
sys.stdin.read().splitlines(), source="<stdin>"
|
||||
)
|
||||
else:
|
||||
path = Path(source)
|
||||
if not path.is_file():
|
||||
raise UsageError(f"{path} not found")
|
||||
references, problems = read_pack_references(path)
|
||||
|
||||
for problem in problems:
|
||||
progress.warn(problem)
|
||||
if not references:
|
||||
raise UsageError(f"no valid sticker pack references in {source}")
|
||||
return references
|
||||
|
||||
|
||||
def collect_interactively(
|
||||
*,
|
||||
types_given: str | None,
|
||||
url_file: Path = DEFAULT_URL_FILE,
|
||||
ask: Callable[[str], str] = input,
|
||||
say: Callable[[str], None] = print,
|
||||
) -> tuple[str | None, list[str]]:
|
||||
"""Ask for the file types and pack references.
|
||||
|
||||
Returns ``(types_text, references)``; ``types_text`` is ``None`` when the
|
||||
caller already supplied ``--types``.
|
||||
"""
|
||||
say(f"--- Telegram Sticker Pack Downloader {__version__} ---")
|
||||
|
||||
if types_given is None:
|
||||
say("")
|
||||
say("Available file types:")
|
||||
for name in ALL_FILE_TYPES:
|
||||
say(f" {name:<5} {FILE_TYPE_HELP[name]}")
|
||||
types_given = ask("File types (comma-separated, or 'all') [all]: ").strip()
|
||||
|
||||
# Validate now so a typo does not surface after the bot has connected.
|
||||
parse_file_types(types_given)
|
||||
|
||||
say("")
|
||||
say("Choose an option:")
|
||||
say(" 1: download one or more packs from URLs you type")
|
||||
say(f" 2: download every pack listed in {url_file}")
|
||||
choice = ask("Your choice (1 or 2) [1]: ").strip() or "1"
|
||||
|
||||
if choice == "1":
|
||||
answer = ask("Sticker pack URL(s), separated by spaces: ").strip()
|
||||
references = answer.split()
|
||||
if not references:
|
||||
raise UsageError("no sticker pack URL given")
|
||||
return types_given, references
|
||||
|
||||
if choice == "2":
|
||||
if not url_file.is_file():
|
||||
raise UsageError(f"{url_file} not found")
|
||||
references, problems = read_pack_references(url_file)
|
||||
for problem in problems:
|
||||
say(f"warning: {problem}")
|
||||
if not references:
|
||||
raise UsageError(f"no valid sticker pack references in {url_file}")
|
||||
say(f"Found {len(references)} pack(s) in {url_file}.")
|
||||
return types_given, references
|
||||
|
||||
raise UsageError(f"invalid choice {choice!r}; expected 1 or 2")
|
||||
|
||||
|
||||
def _make_progress(args: argparse.Namespace) -> ProgressReporter:
|
||||
if args.quiet:
|
||||
return NullProgress()
|
||||
return ConsoleProgress(sys.stderr, verbose=args.verbose)
|
||||
|
||||
|
||||
def _resolve_token(args: argparse.Namespace) -> str:
|
||||
env_file = Path(args.env_file)
|
||||
if env_file.is_file():
|
||||
load_dotenv(env_file)
|
||||
else:
|
||||
load_dotenv()
|
||||
token = (args.token or os.getenv(TOKEN_ENV_VAR) or "").strip()
|
||||
if not token:
|
||||
raise UsageError(_TOKEN_MISSING)
|
||||
return token
|
||||
|
||||
|
||||
async def _download(
|
||||
token: str,
|
||||
references: Sequence[str],
|
||||
config: DownloadConfig,
|
||||
progress: ProgressReporter,
|
||||
) -> list[PackResult]:
|
||||
timeout = aiohttp.ClientTimeout(total=300, sock_connect=30, sock_read=60)
|
||||
async with Bot(token) as bot:
|
||||
me = await bot.get_me()
|
||||
progress.info(f"Connected as @{me.username}.")
|
||||
async with aiohttp.ClientSession(timeout=timeout) as session:
|
||||
downloader = StickerDownloader(
|
||||
bot=bot,
|
||||
fetcher=AiohttpFetcher(session),
|
||||
config=config,
|
||||
progress=progress,
|
||||
)
|
||||
return await downloader.download_all(references)
|
||||
|
||||
|
||||
def main(argv: Sequence[str] | None = None) -> int:
|
||||
"""Run the CLI. Returns the process exit code."""
|
||||
args = build_parser().parse_args(argv)
|
||||
progress = _make_progress(args)
|
||||
|
||||
try:
|
||||
references = _resolve_references(args, progress)
|
||||
types_text = args.types
|
||||
|
||||
if not references:
|
||||
if not sys.stdin.isatty():
|
||||
raise UsageError(
|
||||
"no packs given. Pass PACK arguments, --from-file PATH, "
|
||||
"or run interactively from a terminal."
|
||||
)
|
||||
types_text, references = collect_interactively(types_given=args.types)
|
||||
|
||||
config = DownloadConfig(
|
||||
file_types=parse_file_types(types_text),
|
||||
output_dir=args.output,
|
||||
concurrency=args.concurrency,
|
||||
retries=args.retries,
|
||||
overwrite=args.overwrite,
|
||||
write_metadata=args.write_metadata,
|
||||
pack_thumbnail=args.pack_thumbnail,
|
||||
emoji_names=args.emoji_names,
|
||||
dry_run=args.dry_run,
|
||||
)
|
||||
token = _resolve_token(args)
|
||||
except (UsageError, StickerDownloaderError, ValueError) as error:
|
||||
print(f"error: {error}", file=sys.stderr)
|
||||
return EXIT_USAGE
|
||||
except (EOFError, KeyboardInterrupt):
|
||||
print("\nAborted.", file=sys.stderr)
|
||||
return EXIT_USAGE
|
||||
|
||||
if config.dry_run:
|
||||
progress.info("Dry run: nothing will be written to disk.")
|
||||
|
||||
try:
|
||||
results = asyncio.run(_download(token, references, config, progress))
|
||||
except InvalidToken:
|
||||
print(f"error: {_TOKEN_REJECTED}", file=sys.stderr)
|
||||
return EXIT_USAGE
|
||||
except TelegramError as error:
|
||||
print(f"error: could not reach Telegram: {error}", file=sys.stderr)
|
||||
return EXIT_FAILURES
|
||||
except KeyboardInterrupt:
|
||||
print("\nInterrupted.", file=sys.stderr)
|
||||
return EXIT_FAILURES
|
||||
|
||||
return EXIT_OK if all(result.ok for result in results) else EXIT_FAILURES
|
||||
|
||||
|
||||
def run() -> int: # pragma: no cover - thin console-script wrapper
|
||||
return main()
|
||||
|
||||
|
||||
if __name__ == "__main__": # pragma: no cover
|
||||
raise SystemExit(main())
|
||||
Reference in New Issue
Block a user