360 lines
13 KiB
Python
360 lines
13 KiB
Python
"""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__
|