big upgrade UwU
This commit is contained in:
@@ -0,0 +1,27 @@
|
||||
"""Download Telegram sticker packs from the command line."""
|
||||
|
||||
from sticker_downloader.config import ALL_FILE_TYPES, DownloadConfig, parse_file_types
|
||||
from sticker_downloader.downloader import FileOutcome, PackResult, StickerDownloader
|
||||
from sticker_downloader.errors import (
|
||||
FetchError,
|
||||
InvalidPackReference,
|
||||
StickerDownloaderError,
|
||||
)
|
||||
from sticker_downloader.urls import parse_pack_name, read_pack_references
|
||||
|
||||
__version__ = "1.0.0"
|
||||
|
||||
__all__ = [
|
||||
"ALL_FILE_TYPES",
|
||||
"DownloadConfig",
|
||||
"FetchError",
|
||||
"FileOutcome",
|
||||
"InvalidPackReference",
|
||||
"PackResult",
|
||||
"StickerDownloader",
|
||||
"StickerDownloaderError",
|
||||
"__version__",
|
||||
"parse_file_types",
|
||||
"parse_pack_name",
|
||||
"read_pack_references",
|
||||
]
|
||||
@@ -0,0 +1,6 @@
|
||||
"""Allow ``python -m sticker_downloader``."""
|
||||
|
||||
from sticker_downloader.cli import main
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -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())
|
||||
@@ -0,0 +1,88 @@
|
||||
"""Download settings and file-type handling."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
|
||||
from sticker_downloader.errors import InvalidFileType
|
||||
|
||||
STATIC = "webp"
|
||||
ANIMATED = "tgs"
|
||||
VIDEO = "webm"
|
||||
THUMBNAIL = "png"
|
||||
|
||||
#: Every type the downloader knows how to save, in a stable display order.
|
||||
ALL_FILE_TYPES: tuple[str, ...] = (STATIC, ANIMATED, VIDEO, THUMBNAIL)
|
||||
|
||||
FILE_TYPE_HELP: dict[str, str] = {
|
||||
STATIC: "static stickers",
|
||||
ANIMATED: "animated stickers (Lottie)",
|
||||
VIDEO: "video stickers",
|
||||
THUMBNAIL: "PNG thumbnails",
|
||||
}
|
||||
|
||||
DEFAULT_OUTPUT_DIR = Path("downloads")
|
||||
DEFAULT_URL_FILE = Path("urls.txt")
|
||||
DEFAULT_CONCURRENCY = 8
|
||||
DEFAULT_RETRIES = 3
|
||||
METADATA_FILE_NAME = "pack.json"
|
||||
PACK_THUMBNAIL_STEM = "_pack_thumbnail"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class DownloadConfig:
|
||||
"""Everything that shapes a download run."""
|
||||
|
||||
file_types: frozenset[str] = field(default_factory=lambda: frozenset(ALL_FILE_TYPES))
|
||||
output_dir: Path = DEFAULT_OUTPUT_DIR
|
||||
concurrency: int = DEFAULT_CONCURRENCY
|
||||
retries: int = DEFAULT_RETRIES
|
||||
overwrite: bool = False
|
||||
write_metadata: bool = True
|
||||
pack_thumbnail: bool = True
|
||||
emoji_names: bool = False
|
||||
dry_run: bool = False
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
unknown = sorted(self.file_types - set(ALL_FILE_TYPES))
|
||||
if unknown:
|
||||
raise InvalidFileType(f"unknown file type(s): {', '.join(unknown)}")
|
||||
if not self.file_types:
|
||||
raise InvalidFileType("at least one file type is required")
|
||||
if self.concurrency < 1:
|
||||
raise ValueError("concurrency must be at least 1")
|
||||
if self.retries < 1:
|
||||
raise ValueError("retries must be at least 1")
|
||||
|
||||
@property
|
||||
def sorted_file_types(self) -> list[str]:
|
||||
"""The requested types in canonical order, for stable output."""
|
||||
return [file_type for file_type in ALL_FILE_TYPES if file_type in self.file_types]
|
||||
|
||||
def wants(self, file_type: str) -> bool:
|
||||
return file_type in self.file_types
|
||||
|
||||
|
||||
def parse_file_types(raw: str | None) -> frozenset[str]:
|
||||
"""Turn user input such as ``"webp, png"``, ``"all"`` or ``""`` into a type set.
|
||||
|
||||
Empty input and ``"all"`` both mean "every supported type", matching what the
|
||||
interactive prompt advertises.
|
||||
"""
|
||||
text = (raw or "").strip().lower()
|
||||
if not text or text == "all":
|
||||
return frozenset(ALL_FILE_TYPES)
|
||||
|
||||
requested = [part.strip().lstrip(".") for part in text.replace(" ", ",").split(",")]
|
||||
selected = {part for part in requested if part}
|
||||
if not selected:
|
||||
raise InvalidFileType("no file types given")
|
||||
|
||||
unknown = sorted(selected - set(ALL_FILE_TYPES))
|
||||
if unknown:
|
||||
raise InvalidFileType(
|
||||
f"unknown file type(s): {', '.join(unknown)}. "
|
||||
f"Supported: {', '.join(ALL_FILE_TYPES)}, or 'all'."
|
||||
)
|
||||
return frozenset(selected)
|
||||
@@ -0,0 +1,359 @@
|
||||
"""The async sticker-pack downloader."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
from collections.abc import Awaitable, Callable, Iterable, Sequence
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from typing import Any, TypeVar
|
||||
|
||||
from telegram.error import BadRequest, TelegramError
|
||||
|
||||
from sticker_downloader.config import (
|
||||
ALL_FILE_TYPES,
|
||||
ANIMATED,
|
||||
METADATA_FILE_NAME,
|
||||
PACK_THUMBNAIL_STEM,
|
||||
STATIC,
|
||||
THUMBNAIL,
|
||||
VIDEO,
|
||||
DownloadConfig,
|
||||
)
|
||||
from sticker_downloader.errors import InvalidPackReference, StickerDownloaderError
|
||||
from sticker_downloader.fetcher import Fetcher
|
||||
from sticker_downloader.progress import NullProgress, ProgressReporter
|
||||
from sticker_downloader.results import FileOutcome, FileStatus, PackResult
|
||||
from sticker_downloader.retry import with_retries
|
||||
from sticker_downloader.urls import parse_pack_name
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
_UNSAFE_FILENAME_CHARS = re.compile(r'[<>:"/\\|?*\x00-\x1f]')
|
||||
_PART_SUFFIX = ".part"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _FileJob:
|
||||
"""One file to put on disk."""
|
||||
|
||||
file_id: str
|
||||
stem: Path
|
||||
#: ``None`` means "derive the extension from Telegram's own file path",
|
||||
#: which is how pack thumbnails are handled (they can be webp, tgs or webm).
|
||||
suffix: str | None
|
||||
|
||||
|
||||
class StickerDownloader:
|
||||
"""Downloads sticker packs concurrently into per-pack directories."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
bot: Any,
|
||||
fetcher: Fetcher,
|
||||
config: DownloadConfig | None = None,
|
||||
progress: ProgressReporter | None = None,
|
||||
*,
|
||||
sleep: Callable[[float], Awaitable[None]] = asyncio.sleep,
|
||||
now: Callable[[], datetime] | None = None,
|
||||
) -> None:
|
||||
self._bot = bot
|
||||
self._fetcher = fetcher
|
||||
self._config = config or DownloadConfig()
|
||||
self._progress: ProgressReporter = progress or NullProgress()
|
||||
self._sleep = sleep
|
||||
self._now = now or (lambda: datetime.now(timezone.utc))
|
||||
|
||||
@property
|
||||
def config(self) -> DownloadConfig:
|
||||
return self._config
|
||||
|
||||
# -- public API -------------------------------------------------------
|
||||
|
||||
async def download_all(self, references: Iterable[str]) -> list[PackResult]:
|
||||
"""Download every reference in turn.
|
||||
|
||||
Packs run one after another so progress stays readable; the files inside a
|
||||
pack are downloaded concurrently.
|
||||
"""
|
||||
results: list[PackResult] = []
|
||||
for reference in references:
|
||||
results.append(await self.download_pack(reference))
|
||||
self._progress.run_finished(results)
|
||||
return results
|
||||
|
||||
async def download_pack(self, reference: str) -> PackResult:
|
||||
"""Download a single pack, never raising for an expected failure."""
|
||||
try:
|
||||
pack_name = parse_pack_name(reference)
|
||||
except InvalidPackReference as error:
|
||||
result = PackResult(reference=reference, error=str(error))
|
||||
self._progress.pack_finished(result)
|
||||
return result
|
||||
|
||||
try:
|
||||
sticker_set = await self._call(lambda: self._bot.get_sticker_set(pack_name))
|
||||
except Exception as error:
|
||||
result = PackResult(
|
||||
reference=reference, name=pack_name, error=_describe(error)
|
||||
)
|
||||
self._progress.pack_finished(result)
|
||||
return result
|
||||
|
||||
stickers: Sequence[Any] = list(getattr(sticker_set, "stickers", None) or [])
|
||||
name = getattr(sticker_set, "name", None) or pack_name
|
||||
directory = self._config.output_dir / name
|
||||
|
||||
jobs = self._plan_jobs(sticker_set, stickers, directory)
|
||||
self._progress.pack_started(reference, name, len(stickers), len(jobs))
|
||||
|
||||
if not self._config.dry_run:
|
||||
await asyncio.to_thread(directory.mkdir, parents=True, exist_ok=True)
|
||||
|
||||
semaphore = asyncio.Semaphore(self._config.concurrency)
|
||||
outcomes = list(
|
||||
await asyncio.gather(*(self._run_job(job, semaphore) for job in jobs))
|
||||
)
|
||||
|
||||
result = PackResult(
|
||||
reference=reference,
|
||||
name=name,
|
||||
title=getattr(sticker_set, "title", None),
|
||||
directory=directory,
|
||||
total_stickers=len(stickers),
|
||||
outcomes=outcomes,
|
||||
)
|
||||
|
||||
if self._config.write_metadata and not self._config.dry_run:
|
||||
try:
|
||||
await self._write_metadata(sticker_set, stickers, directory, result)
|
||||
except OSError as error:
|
||||
self._progress.warn(f"could not write {METADATA_FILE_NAME}: {error}")
|
||||
|
||||
self._progress.pack_finished(result)
|
||||
return result
|
||||
|
||||
# -- planning ---------------------------------------------------------
|
||||
|
||||
def _plan_jobs(
|
||||
self, sticker_set: Any, stickers: Sequence[Any], directory: Path
|
||||
) -> list[_FileJob]:
|
||||
jobs: list[_FileJob] = []
|
||||
|
||||
set_thumbnail = getattr(sticker_set, "thumbnail", None)
|
||||
if self._config.pack_thumbnail and set_thumbnail is not None:
|
||||
jobs.append(
|
||||
_FileJob(
|
||||
file_id=set_thumbnail.file_id,
|
||||
stem=directory / PACK_THUMBNAIL_STEM,
|
||||
suffix=None,
|
||||
)
|
||||
)
|
||||
|
||||
for index, sticker in enumerate(stickers, start=1):
|
||||
base = self._sticker_stem(index, sticker)
|
||||
|
||||
primary = primary_file_type(sticker)
|
||||
if self._config.wants(primary):
|
||||
jobs.append(
|
||||
_FileJob(
|
||||
file_id=sticker.file_id,
|
||||
stem=directory / base,
|
||||
suffix=f".{primary}",
|
||||
)
|
||||
)
|
||||
|
||||
thumbnail = getattr(sticker, "thumbnail", None)
|
||||
if self._config.wants(THUMBNAIL) and thumbnail is not None:
|
||||
jobs.append(
|
||||
_FileJob(
|
||||
file_id=thumbnail.file_id,
|
||||
stem=directory / f"{base}_thumb",
|
||||
suffix=f".{THUMBNAIL}",
|
||||
)
|
||||
)
|
||||
|
||||
return jobs
|
||||
|
||||
def _sticker_stem(self, index: int, sticker: Any) -> str:
|
||||
base = f"{index:03d}"
|
||||
if not self._config.emoji_names:
|
||||
return base
|
||||
emoji = sanitize_filename_part(getattr(sticker, "emoji", None) or "")
|
||||
return f"{base}_{emoji}" if emoji else base
|
||||
|
||||
# -- execution --------------------------------------------------------
|
||||
|
||||
async def _run_job(self, job: _FileJob, semaphore: asyncio.Semaphore) -> FileOutcome:
|
||||
async with semaphore:
|
||||
outcome = await self._process_job(job)
|
||||
self._progress.file_finished(outcome)
|
||||
return outcome
|
||||
|
||||
async def _process_job(self, job: _FileJob) -> FileOutcome:
|
||||
planned_path = self._planned_path(job)
|
||||
try:
|
||||
existing = None if self._config.overwrite else self._existing_path(job)
|
||||
if existing is not None:
|
||||
return FileOutcome(existing, FileStatus.SKIPPED_EXISTING)
|
||||
|
||||
if self._config.dry_run:
|
||||
return FileOutcome(planned_path, FileStatus.PLANNED)
|
||||
|
||||
file = await self._call(lambda: self._bot.get_file(job.file_id))
|
||||
remote_path = getattr(file, "file_path", None)
|
||||
if not remote_path:
|
||||
raise StickerDownloaderError("Telegram returned no file path")
|
||||
|
||||
destination = self._resolve_path(job, remote_path)
|
||||
data = await self._call(lambda: self._fetcher.fetch(remote_path))
|
||||
await asyncio.to_thread(_write_atomically, destination, data)
|
||||
return FileOutcome(destination, FileStatus.DOWNLOADED, size=len(data))
|
||||
except Exception as error:
|
||||
return FileOutcome(planned_path, FileStatus.FAILED, error=_describe(error))
|
||||
|
||||
async def _call(self, operation: Callable[[], Awaitable[T]]) -> T:
|
||||
"""Run one network operation with this run's retry policy."""
|
||||
return await with_retries(
|
||||
operation,
|
||||
attempts=self._config.retries,
|
||||
on_retry=self._on_retry,
|
||||
sleep=self._sleep,
|
||||
)
|
||||
|
||||
def _on_retry(self, error: BaseException, attempt: int, delay: float) -> None:
|
||||
self._progress.warn(
|
||||
f"{_describe(error)} — retrying in {delay:.1f}s (attempt {attempt + 1})"
|
||||
)
|
||||
|
||||
# -- paths ------------------------------------------------------------
|
||||
|
||||
def _planned_path(self, job: _FileJob) -> Path:
|
||||
"""Best guess at the destination, used for dry runs and error messages."""
|
||||
return self._resolve_path(job, f"unknown.{STATIC}")
|
||||
|
||||
def _resolve_path(self, job: _FileJob, remote_path: str) -> Path:
|
||||
suffix = job.suffix if job.suffix is not None else _suffix_from_url(remote_path)
|
||||
return job.stem.parent / f"{job.stem.name}{suffix}"
|
||||
|
||||
def _existing_path(self, job: _FileJob) -> Path | None:
|
||||
"""The already-downloaded file for ``job``, if there is one.
|
||||
|
||||
Checking before calling ``get_file`` means resuming a large batch costs no
|
||||
API calls for the parts that are already on disk.
|
||||
"""
|
||||
if job.suffix is not None:
|
||||
candidate = job.stem.parent / f"{job.stem.name}{job.suffix}"
|
||||
return candidate if candidate.is_file() else None
|
||||
for file_type in ALL_FILE_TYPES:
|
||||
candidate = job.stem.parent / f"{job.stem.name}.{file_type}"
|
||||
if candidate.is_file():
|
||||
return candidate
|
||||
return None
|
||||
|
||||
# -- metadata ---------------------------------------------------------
|
||||
|
||||
async def _write_metadata(
|
||||
self,
|
||||
sticker_set: Any,
|
||||
stickers: Sequence[Any],
|
||||
directory: Path,
|
||||
result: PackResult,
|
||||
) -> None:
|
||||
present = {
|
||||
outcome.path
|
||||
for outcome in result.outcomes
|
||||
if outcome.status in (FileStatus.DOWNLOADED, FileStatus.SKIPPED_EXISTING)
|
||||
}
|
||||
|
||||
entries = []
|
||||
for index, sticker in enumerate(stickers, start=1):
|
||||
base = self._sticker_stem(index, sticker)
|
||||
files = sorted(
|
||||
path.name
|
||||
for path in present
|
||||
if path.stem == base or path.stem == f"{base}_thumb"
|
||||
)
|
||||
entries.append(
|
||||
{
|
||||
"index": index,
|
||||
"emoji": getattr(sticker, "emoji", None),
|
||||
"file_unique_id": getattr(sticker, "file_unique_id", None),
|
||||
"type": _enum_value(getattr(sticker, "type", None)),
|
||||
"width": getattr(sticker, "width", None),
|
||||
"height": getattr(sticker, "height", None),
|
||||
"is_animated": bool(getattr(sticker, "is_animated", False)),
|
||||
"is_video": bool(getattr(sticker, "is_video", False)),
|
||||
"files": files,
|
||||
}
|
||||
)
|
||||
|
||||
from sticker_downloader import __version__
|
||||
|
||||
payload = {
|
||||
"pack": {
|
||||
"reference": result.reference,
|
||||
"name": result.name,
|
||||
"title": result.title,
|
||||
"sticker_type": _enum_value(getattr(sticker_set, "sticker_type", None)),
|
||||
},
|
||||
"generated_at": self._now().isoformat(),
|
||||
"downloader_version": __version__,
|
||||
"requested_file_types": self._config.sorted_file_types,
|
||||
"sticker_count": len(stickers),
|
||||
"stickers": entries,
|
||||
}
|
||||
|
||||
text = json.dumps(payload, indent=2, ensure_ascii=False) + "\n"
|
||||
await asyncio.to_thread(
|
||||
(directory / METADATA_FILE_NAME).write_text, text, encoding="utf-8"
|
||||
)
|
||||
|
||||
|
||||
# -- module-level helpers -------------------------------------------------
|
||||
|
||||
|
||||
def primary_file_type(sticker: Any) -> str:
|
||||
"""The file type Telegram stores this sticker as."""
|
||||
if getattr(sticker, "is_animated", False):
|
||||
return ANIMATED
|
||||
if getattr(sticker, "is_video", False):
|
||||
return VIDEO
|
||||
return STATIC
|
||||
|
||||
|
||||
def sanitize_filename_part(text: str) -> str:
|
||||
"""Strip characters that are illegal or awkward inside a file name."""
|
||||
cleaned = _UNSAFE_FILENAME_CHARS.sub("", text).strip(" .")
|
||||
return cleaned[:32]
|
||||
|
||||
|
||||
def _suffix_from_url(remote_path: str) -> str:
|
||||
suffix = Path(remote_path.split("?")[0]).suffix.lower()
|
||||
return suffix if suffix else f".{STATIC}"
|
||||
|
||||
|
||||
def _write_atomically(destination: Path, data: bytes) -> None:
|
||||
"""Write via a temporary file so an interrupted run leaves no partial file."""
|
||||
destination.parent.mkdir(parents=True, exist_ok=True)
|
||||
temporary = destination.with_name(destination.name + _PART_SUFFIX)
|
||||
temporary.write_bytes(data)
|
||||
os.replace(temporary, destination)
|
||||
|
||||
|
||||
def _enum_value(value: Any) -> Any:
|
||||
return getattr(value, "value", value)
|
||||
|
||||
|
||||
def _describe(error: BaseException) -> str:
|
||||
"""A short, user-facing description of a failure."""
|
||||
if isinstance(error, BadRequest) and "stickerset_invalid" in str(error).lower():
|
||||
return "sticker pack not found"
|
||||
if isinstance(error, TelegramError):
|
||||
return f"{type(error).__name__}: {error}"
|
||||
message = str(error).strip()
|
||||
return message or type(error).__name__
|
||||
@@ -0,0 +1,31 @@
|
||||
"""Exception types used across the package."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
|
||||
class StickerDownloaderError(Exception):
|
||||
"""Base class for every error raised by this package."""
|
||||
|
||||
|
||||
class InvalidPackReference(StickerDownloaderError):
|
||||
"""Raised when a string cannot be read as a sticker pack reference."""
|
||||
|
||||
|
||||
class InvalidFileType(StickerDownloaderError):
|
||||
"""Raised when the user asks for a file type we cannot download."""
|
||||
|
||||
|
||||
class FetchError(StickerDownloaderError):
|
||||
"""Raised when the CDN returns a non-success status for a file."""
|
||||
|
||||
def __init__(self, url: str, status: int, reason: str = "") -> None:
|
||||
detail = f" ({reason})" if reason else ""
|
||||
super().__init__(f"HTTP {status}{detail} for {url}")
|
||||
self.url = url
|
||||
self.status = status
|
||||
self.reason = reason
|
||||
|
||||
@property
|
||||
def retriable(self) -> bool:
|
||||
"""Rate limits and server-side faults are worth another attempt."""
|
||||
return self.status == 429 or self.status >= 500
|
||||
@@ -0,0 +1,29 @@
|
||||
"""Fetching file bytes from Telegram's CDN."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Protocol, runtime_checkable
|
||||
|
||||
import aiohttp
|
||||
|
||||
from sticker_downloader.errors import FetchError
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class Fetcher(Protocol):
|
||||
"""Minimal interface the downloader needs to read a file off the network."""
|
||||
|
||||
async def fetch(self, url: str) -> bytes: ...
|
||||
|
||||
|
||||
class AiohttpFetcher:
|
||||
"""A :class:`Fetcher` backed by a shared :class:`aiohttp.ClientSession`."""
|
||||
|
||||
def __init__(self, session: aiohttp.ClientSession) -> None:
|
||||
self._session = session
|
||||
|
||||
async def fetch(self, url: str) -> bytes:
|
||||
async with self._session.get(url) as response:
|
||||
if response.status >= 400:
|
||||
raise FetchError(url, response.status, response.reason or "")
|
||||
return await response.read()
|
||||
@@ -0,0 +1,158 @@
|
||||
"""Console reporting for a download run."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
from typing import IO, Protocol
|
||||
|
||||
from sticker_downloader.results import (
|
||||
FileOutcome,
|
||||
FileStatus,
|
||||
PackResult,
|
||||
RunTotals,
|
||||
format_size,
|
||||
)
|
||||
|
||||
|
||||
class ProgressReporter(Protocol):
|
||||
"""Everything the downloader and CLI report as a run unfolds."""
|
||||
|
||||
def info(self, message: str) -> None: ...
|
||||
def warn(self, message: str) -> None: ...
|
||||
def error(self, message: str) -> None: ...
|
||||
def pack_started(
|
||||
self, reference: str, name: str, total_stickers: int, total_files: int
|
||||
) -> None: ...
|
||||
def file_finished(self, outcome: FileOutcome) -> None: ...
|
||||
def pack_finished(self, result: PackResult) -> None: ...
|
||||
def run_finished(self, results: list[PackResult]) -> None: ...
|
||||
|
||||
|
||||
class NullProgress:
|
||||
"""Reports nothing. Used by ``--quiet`` and by tests."""
|
||||
|
||||
def info(self, message: str) -> None: ...
|
||||
def warn(self, message: str) -> None: ...
|
||||
def error(self, message: str) -> None: ...
|
||||
def pack_started(
|
||||
self, reference: str, name: str, total_stickers: int, total_files: int
|
||||
) -> None: ...
|
||||
def file_finished(self, outcome: FileOutcome) -> None: ...
|
||||
def pack_finished(self, result: PackResult) -> None: ...
|
||||
def run_finished(self, results: list[PackResult]) -> None: ...
|
||||
|
||||
|
||||
class ConsoleProgress:
|
||||
"""Writes human-friendly progress to a stream.
|
||||
|
||||
On a terminal the per-pack counters are redrawn in place; when the output is
|
||||
piped to a file each pack reports a single summary line instead, so logs stay
|
||||
readable.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
stream: IO[str] | None = None,
|
||||
*,
|
||||
verbose: bool = False,
|
||||
use_ansi: bool | None = None,
|
||||
) -> None:
|
||||
self._stream = stream if stream is not None else sys.stderr
|
||||
self._verbose = verbose
|
||||
if use_ansi is None:
|
||||
use_ansi = bool(getattr(self._stream, "isatty", lambda: False)())
|
||||
self._use_ansi = use_ansi
|
||||
self._live_width = 0
|
||||
self._pack_total = 0
|
||||
self._pack_done = 0
|
||||
|
||||
# -- plain messages ---------------------------------------------------
|
||||
|
||||
def info(self, message: str) -> None:
|
||||
self._write_line(message)
|
||||
|
||||
def warn(self, message: str) -> None:
|
||||
self._write_line(f"warning: {message}")
|
||||
|
||||
def error(self, message: str) -> None:
|
||||
self._write_line(f"error: {message}")
|
||||
|
||||
# -- run lifecycle ----------------------------------------------------
|
||||
|
||||
def pack_started(
|
||||
self, reference: str, name: str, total_stickers: int, total_files: int
|
||||
) -> None:
|
||||
self._pack_total = total_files
|
||||
self._pack_done = 0
|
||||
self._write_line(f"\n{name}: {total_stickers} stickers, {total_files} files")
|
||||
|
||||
def file_finished(self, outcome: FileOutcome) -> None:
|
||||
self._pack_done += 1
|
||||
if outcome.status is FileStatus.FAILED:
|
||||
self.error(f"{outcome.name}: {outcome.error}")
|
||||
return
|
||||
if self._verbose:
|
||||
detail = (
|
||||
f" ({format_size(outcome.size)})"
|
||||
if outcome.status is FileStatus.DOWNLOADED
|
||||
else ""
|
||||
)
|
||||
self._write_line(f" {outcome.status.value}: {outcome.name}{detail}")
|
||||
return
|
||||
total = self._pack_total or self._pack_done
|
||||
self._draw_live(f" {self._pack_done}/{total} files")
|
||||
|
||||
def pack_finished(self, result: PackResult) -> None:
|
||||
self._clear_live()
|
||||
if result.error is not None:
|
||||
self.error(f"{result.reference}: {result.error}")
|
||||
return
|
||||
location = f" -> {result.directory}" if result.directory else ""
|
||||
self._write_line(f" {result.summary()}{location}")
|
||||
|
||||
def run_finished(self, results: list[PackResult]) -> None:
|
||||
self._clear_live()
|
||||
totals = RunTotals.from_results(results)
|
||||
if totals.packs == 0:
|
||||
self._write_line("\nNothing to do.")
|
||||
return
|
||||
written = format_size(totals.bytes_written)
|
||||
lines = [
|
||||
"",
|
||||
f"Done: {totals.packs_ok}/{totals.packs} packs",
|
||||
f" files downloaded : {totals.downloaded} ({written})",
|
||||
]
|
||||
if totals.planned:
|
||||
lines.append(f" files planned : {totals.planned}")
|
||||
if totals.skipped:
|
||||
lines.append(f" already present : {totals.skipped}")
|
||||
if totals.failed:
|
||||
lines.append(f" files failed : {totals.failed}")
|
||||
if totals.packs_failed:
|
||||
lines.append(f" packs failed : {totals.packs_failed}")
|
||||
for result in results:
|
||||
if not result.ok:
|
||||
reason = result.error or f"{result.failed} file(s) failed"
|
||||
lines.append(f" - {result.reference}: {reason}")
|
||||
self._write_line("\n".join(lines))
|
||||
|
||||
# -- internals --------------------------------------------------------
|
||||
|
||||
def _draw_live(self, text: str) -> None:
|
||||
if not self._use_ansi:
|
||||
return
|
||||
padding = " " * max(0, self._live_width - len(text))
|
||||
self._stream.write(f"\r{text}{padding}")
|
||||
self._stream.flush()
|
||||
self._live_width = len(text)
|
||||
|
||||
def _clear_live(self) -> None:
|
||||
if self._use_ansi and self._live_width:
|
||||
self._stream.write("\r" + " " * self._live_width + "\r")
|
||||
self._stream.flush()
|
||||
self._live_width = 0
|
||||
|
||||
def _write_line(self, text: str) -> None:
|
||||
self._clear_live()
|
||||
self._stream.write(f"{text}\n")
|
||||
self._stream.flush()
|
||||
@@ -0,0 +1,118 @@
|
||||
"""Result objects describing what a download run did."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
class FileStatus(str, Enum):
|
||||
"""Outcome of a single file within a pack."""
|
||||
|
||||
DOWNLOADED = "downloaded"
|
||||
SKIPPED_EXISTING = "skipped_existing"
|
||||
FAILED = "failed"
|
||||
PLANNED = "planned" # --dry-run only
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class FileOutcome:
|
||||
"""What happened to one sticker file."""
|
||||
|
||||
path: Path
|
||||
status: FileStatus
|
||||
size: int = 0
|
||||
error: str | None = None
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return self.path.name
|
||||
|
||||
|
||||
@dataclass
|
||||
class PackResult:
|
||||
"""Aggregated outcome for one sticker pack."""
|
||||
|
||||
reference: str
|
||||
name: str | None = None
|
||||
title: str | None = None
|
||||
directory: Path | None = None
|
||||
total_stickers: int = 0
|
||||
outcomes: list[FileOutcome] = field(default_factory=list)
|
||||
error: str | None = None
|
||||
|
||||
@property
|
||||
def downloaded(self) -> int:
|
||||
return sum(1 for o in self.outcomes if o.status is FileStatus.DOWNLOADED)
|
||||
|
||||
@property
|
||||
def skipped(self) -> int:
|
||||
return sum(1 for o in self.outcomes if o.status is FileStatus.SKIPPED_EXISTING)
|
||||
|
||||
@property
|
||||
def planned(self) -> int:
|
||||
return sum(1 for o in self.outcomes if o.status is FileStatus.PLANNED)
|
||||
|
||||
@property
|
||||
def failed(self) -> int:
|
||||
return sum(1 for o in self.outcomes if o.status is FileStatus.FAILED)
|
||||
|
||||
@property
|
||||
def bytes_written(self) -> int:
|
||||
return sum(o.size for o in self.outcomes if o.status is FileStatus.DOWNLOADED)
|
||||
|
||||
@property
|
||||
def ok(self) -> bool:
|
||||
"""True when the pack was processed without a fatal or per-file error."""
|
||||
return self.error is None and self.failed == 0
|
||||
|
||||
def summary(self) -> str:
|
||||
if self.error is not None:
|
||||
return f"failed: {self.error}"
|
||||
parts = [f"{self.downloaded} downloaded"]
|
||||
if self.planned:
|
||||
parts.append(f"{self.planned} planned")
|
||||
if self.skipped:
|
||||
parts.append(f"{self.skipped} already present")
|
||||
if self.failed:
|
||||
parts.append(f"{self.failed} failed")
|
||||
return ", ".join(parts)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class RunTotals:
|
||||
"""Totals across every pack in a run."""
|
||||
|
||||
packs: int
|
||||
packs_ok: int
|
||||
packs_failed: int
|
||||
downloaded: int
|
||||
skipped: int
|
||||
planned: int
|
||||
failed: int
|
||||
bytes_written: int
|
||||
|
||||
@classmethod
|
||||
def from_results(cls, results: list[PackResult]) -> RunTotals:
|
||||
return cls(
|
||||
packs=len(results),
|
||||
packs_ok=sum(1 for r in results if r.ok),
|
||||
packs_failed=sum(1 for r in results if not r.ok),
|
||||
downloaded=sum(r.downloaded for r in results),
|
||||
skipped=sum(r.skipped for r in results),
|
||||
planned=sum(r.planned for r in results),
|
||||
failed=sum(r.failed for r in results),
|
||||
bytes_written=sum(r.bytes_written for r in results),
|
||||
)
|
||||
|
||||
|
||||
def format_size(num_bytes: int) -> str:
|
||||
"""Human-readable byte count, e.g. ``1.4 MiB``."""
|
||||
size = float(num_bytes)
|
||||
for unit in ("B", "KiB", "MiB", "GiB"):
|
||||
if size < 1024 or unit == "GiB":
|
||||
precision = 0 if unit == "B" else 1
|
||||
return f"{size:.{precision}f} {unit}"
|
||||
size /= 1024
|
||||
raise AssertionError("unreachable") # pragma: no cover
|
||||
@@ -0,0 +1,84 @@
|
||||
"""Retry helper shared by the Telegram API calls and the file downloads."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from collections.abc import Awaitable, Callable
|
||||
from typing import TypeVar
|
||||
|
||||
import aiohttp
|
||||
from telegram.error import BadRequest, Forbidden, InvalidToken, NetworkError, RetryAfter
|
||||
|
||||
from sticker_downloader.errors import FetchError
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
#: Transient failures. ``BadRequest`` subclasses ``NetworkError`` in
|
||||
#: python-telegram-bot, so it has to be excluded explicitly below — retrying a
|
||||
#: "sticker set is invalid" response would just waste three round-trips.
|
||||
_RETRIABLE = (NetworkError, aiohttp.ClientError, asyncio.TimeoutError, OSError)
|
||||
_NEVER_RETRIABLE = (BadRequest, Forbidden, InvalidToken)
|
||||
|
||||
|
||||
def is_retriable(error: BaseException) -> bool:
|
||||
"""Whether another attempt could plausibly succeed."""
|
||||
if isinstance(error, RetryAfter):
|
||||
return True
|
||||
if isinstance(error, _NEVER_RETRIABLE):
|
||||
return False
|
||||
if isinstance(error, FetchError):
|
||||
return error.retriable
|
||||
return isinstance(error, _RETRIABLE)
|
||||
|
||||
|
||||
def retry_delay(error: BaseException, attempt: int, base_delay: float) -> float:
|
||||
"""Seconds to wait before ``attempt`` (1-based) is retried.
|
||||
|
||||
Honours Telegram's own ``retry_after`` hint when it sends one, and otherwise
|
||||
backs off exponentially.
|
||||
"""
|
||||
if isinstance(error, RetryAfter):
|
||||
retry_after = getattr(error, "retry_after", None)
|
||||
if isinstance(retry_after, (int, float)):
|
||||
return float(retry_after)
|
||||
# Newer python-telegram-bot versions may hand back a timedelta.
|
||||
seconds = getattr(retry_after, "total_seconds", None)
|
||||
if callable(seconds):
|
||||
return float(seconds())
|
||||
return base_delay * (2 ** (attempt - 1))
|
||||
|
||||
|
||||
async def with_retries(
|
||||
operation: Callable[[], Awaitable[T]],
|
||||
*,
|
||||
attempts: int = 3,
|
||||
base_delay: float = 0.5,
|
||||
on_retry: Callable[[BaseException, int, float], None] | None = None,
|
||||
sleep: Callable[[float], Awaitable[None]] = asyncio.sleep,
|
||||
) -> T:
|
||||
"""Await ``operation``, retrying transient failures up to ``attempts`` times.
|
||||
|
||||
Args:
|
||||
operation: A zero-argument coroutine function; called afresh per attempt.
|
||||
attempts: Total number of tries, including the first one.
|
||||
base_delay: Seed for the exponential backoff, in seconds.
|
||||
on_retry: Called with ``(error, attempt, delay)`` before each retry.
|
||||
sleep: Injectable sleep, so tests do not have to wait.
|
||||
|
||||
Raises:
|
||||
The last error, once the attempts are exhausted or it is not retriable.
|
||||
"""
|
||||
last_error: BaseException
|
||||
for attempt in range(1, attempts + 1):
|
||||
try:
|
||||
return await operation()
|
||||
except Exception as error:
|
||||
last_error = error
|
||||
if attempt == attempts or not is_retriable(error):
|
||||
raise
|
||||
delay = retry_delay(error, attempt, base_delay)
|
||||
if on_retry is not None:
|
||||
on_retry(error, attempt, delay)
|
||||
await sleep(delay)
|
||||
|
||||
raise last_error # pragma: no cover — loop always returns or raises
|
||||
@@ -0,0 +1,125 @@
|
||||
"""Turn the many ways of naming a sticker pack into a bare pack name."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from collections.abc import Iterable
|
||||
from pathlib import Path
|
||||
from urllib.parse import parse_qs, urlparse
|
||||
|
||||
from sticker_downloader.errors import InvalidPackReference
|
||||
|
||||
#: Telegram pack names are ASCII letters, digits and underscores.
|
||||
_PACK_NAME_RE = re.compile(r"^[A-Za-z0-9_]{1,64}$")
|
||||
|
||||
#: Hosts that serve ``/addstickers/<name>`` links.
|
||||
_KNOWN_HOSTS = frozenset({"t.me", "telegram.me", "telegram.dog", "telesco.pe"})
|
||||
|
||||
_ADD_STICKERS_PATHS = frozenset({"addstickers", "addemoji"})
|
||||
|
||||
|
||||
def parse_pack_name(reference: str) -> str:
|
||||
"""Extract the pack name from anything a user is likely to paste.
|
||||
|
||||
Accepts full URLs (``https://t.me/addstickers/Foo``), scheme-less links
|
||||
(``t.me/addstickers/Foo``), ``tg://addstickers?set=Foo`` deep links, a bare
|
||||
``addstickers/Foo`` path, and a plain pack name.
|
||||
|
||||
Raises:
|
||||
InvalidPackReference: if no pack name can be recovered.
|
||||
"""
|
||||
text = (reference or "").strip().strip('"').strip("'")
|
||||
if not text:
|
||||
raise InvalidPackReference("empty sticker pack reference")
|
||||
|
||||
if text.lower().startswith("tg://"):
|
||||
return _validate(_name_from_deep_link(text), reference)
|
||||
|
||||
if "/" not in text:
|
||||
return _validate(text, reference)
|
||||
|
||||
return _validate(_name_from_url(text, reference), reference)
|
||||
|
||||
|
||||
def _name_from_deep_link(text: str) -> str:
|
||||
parsed = urlparse(text)
|
||||
query = parse_qs(parsed.query)
|
||||
for key in ("set", "name"):
|
||||
values = query.get(key)
|
||||
if values and values[0].strip():
|
||||
return values[0].strip()
|
||||
return ""
|
||||
|
||||
|
||||
def _name_from_url(text: str, reference: str) -> str:
|
||||
candidate = text if "//" in text else f"https://{text}"
|
||||
parsed = urlparse(candidate)
|
||||
segments = [segment for segment in parsed.path.split("/") if segment]
|
||||
|
||||
host = parsed.netloc.lower().removeprefix("www.").split(":")[0]
|
||||
if host in _ADD_STICKERS_PATHS:
|
||||
# A path-only reference such as "addstickers/Foo".
|
||||
return segments[0] if segments else ""
|
||||
|
||||
if host not in _KNOWN_HOSTS:
|
||||
raise InvalidPackReference(
|
||||
f"{reference!r} is not a Telegram sticker link "
|
||||
"(expected something like https://t.me/addstickers/PackName)"
|
||||
)
|
||||
|
||||
if len(segments) < 2 or segments[0].lower() not in _ADD_STICKERS_PATHS:
|
||||
raise InvalidPackReference(f"{reference!r} is not an /addstickers/ link")
|
||||
return segments[1]
|
||||
|
||||
|
||||
def _validate(name: str, reference: str) -> str:
|
||||
name = name.strip()
|
||||
if not _PACK_NAME_RE.match(name):
|
||||
raise InvalidPackReference(
|
||||
f"{reference!r} does not contain a valid pack name "
|
||||
"(letters, digits and underscores only)"
|
||||
)
|
||||
return name
|
||||
|
||||
|
||||
def parse_pack_reference_lines(
|
||||
lines: Iterable[str], source: str = "input"
|
||||
) -> tuple[list[str], list[str]]:
|
||||
"""Parse pack references from lines of text.
|
||||
|
||||
Blank lines and ``#`` comments are ignored. Duplicate references collapse to
|
||||
the first occurrence so a pack is never downloaded twice in one run.
|
||||
|
||||
Returns:
|
||||
``(references, problems)`` — valid references in input order, plus a
|
||||
human-readable message for every line that could not be parsed.
|
||||
"""
|
||||
references: list[str] = []
|
||||
problems: list[str] = []
|
||||
seen: set[str] = set()
|
||||
|
||||
for line_number, raw_line in enumerate(lines, start=1):
|
||||
line = raw_line.strip()
|
||||
if not line or line.startswith("#"):
|
||||
continue
|
||||
try:
|
||||
name = parse_pack_name(line)
|
||||
except InvalidPackReference as error:
|
||||
problems.append(f"{source}:{line_number}: {error}")
|
||||
continue
|
||||
if name.lower() in seen:
|
||||
continue
|
||||
seen.add(name.lower())
|
||||
references.append(line)
|
||||
|
||||
return references, problems
|
||||
|
||||
|
||||
def read_pack_references(path: Path | str) -> tuple[list[str], list[str]]:
|
||||
"""Read pack references from a text file, one per line.
|
||||
|
||||
See :func:`parse_pack_reference_lines` for the accepted syntax.
|
||||
"""
|
||||
file_path = Path(path)
|
||||
lines = file_path.read_text(encoding="utf-8").splitlines()
|
||||
return parse_pack_reference_lines(lines, source=str(file_path))
|
||||
Reference in New Issue
Block a user