"""File I/O utilities: base64 encoding, output path generation, download.""" from __future__ import annotations import base64 from datetime import datetime from pathlib import Path from typing import Optional import httpx from rich.console import Console from rich.progress import ( BarColumn, DownloadColumn, Progress, SpinnerColumn, TimeElapsedColumn, TransferSpeedColumn, ) def image_to_base64(path: Path) -> str: """Read an image file and return a base64-encoded string.""" suffix = path.suffix.lower().lstrip(".") mime_map = { "jpg": "image/jpeg", "jpeg": "image/jpeg", "png": "image/png", "gif": "image/gif", "webp": "image/webp", } mime = mime_map.get(suffix, "image/jpeg") with open(path, "rb") as f: encoded = base64.b64encode(f.read()).decode() return f"data:{mime};base64,{encoded}" def video_to_base64(path: Path) -> str: """Read a video file and return a base64-encoded string.""" suffix = path.suffix.lower().lstrip(".") mime_map = { "mp4": "video/mp4", "mov": "video/quicktime", "avi": "video/x-msvideo", "webm": "video/webm", } mime = mime_map.get(suffix, "video/mp4") with open(path, "rb") as f: encoded = base64.b64encode(f.read()).decode() return f"data:{mime};base64,{encoded}" def auto_output_path(task_type: str, model: str, ext: str = "jpg", output_dir: str = ".") -> Path: """Generate a timestamped output filename.""" ts = datetime.now().strftime("%Y%m%d_%H%M%S") safe_model = model.replace("/", "-").replace(" ", "-") filename = f"freepik_{task_type}_{safe_model}_{ts}.{ext}" return Path(output_dir) / filename def save_from_url(url: str, dest: Path, console: Console) -> None: """Stream-download a URL to a file, displaying a Rich progress bar.""" dest.parent.mkdir(parents=True, exist_ok=True) with Progress( SpinnerColumn(style="bold magenta"), "[progress.description]{task.description}", BarColumn(bar_width=30, style="magenta", complete_style="green"), "[progress.percentage]{task.percentage:>3.0f}%", DownloadColumn(), TransferSpeedColumn(), TimeElapsedColumn(), console=console, transient=True, ) as progress: with httpx.stream("GET", url, follow_redirects=True, timeout=120) as r: r.raise_for_status() total = int(r.headers.get("content-length", 0)) or None task = progress.add_task( f"[dim]Saving[/dim] [bold]{dest.name}[/bold]", total=total, ) with open(dest, "wb") as f: for chunk in r.iter_bytes(chunk_size=65536): f.write(chunk) progress.advance(task, len(chunk)) def get_image_dimensions(path: Path) -> tuple[Optional[int], Optional[int]]: """Return (width, height) of an image using Pillow.""" try: from PIL import Image with Image.open(path) as img: return img.width, img.height except Exception: return None, None