"""Image downloader utility for NFO media files. This module provides functions to download poster, logo, and fanart images from TMDB and validate them. Example: >>> downloader = ImageDownloader() >>> await downloader.download_poster(poster_url, "/path/to/poster.jpg") """ import asyncio import logging from pathlib import Path from typing import Optional import aiohttp from PIL import Image logger = logging.getLogger(__name__) class ImageDownloadError(Exception): """Exception raised for image download failures.""" pass class ImageDownloader: """Utility for downloading and validating images. Supports async context manager protocol for proper resource cleanup. Attributes: max_retries: Maximum retry attempts for downloads timeout: Request timeout in seconds min_file_size: Minimum valid file size in bytes session: Optional aiohttp session (managed internally) Example: >>> async with ImageDownloader() as downloader: ... await downloader.download_poster(url, path) """ def __init__( self, max_retries: int = 3, timeout: int = 30, min_file_size: int = 1024, # 1 KB retry_delay: float = 1.0 ): """Initialize image downloader. Args: max_retries: Maximum retry attempts timeout: Request timeout in seconds min_file_size: Minimum valid file size in bytes retry_delay: Delay between retries in seconds """ self.max_retries = max_retries self.timeout = timeout self.min_file_size = min_file_size self.retry_delay = retry_delay self.session: Optional[aiohttp.ClientSession] = None async def __aenter__(self): """Enter async context manager and create session.""" self._get_session() # Ensure session is created return self async def __aexit__(self, exc_type, exc_val, exc_tb): """Exit async context manager and cleanup resources.""" await self.close() return False async def close(self): """Close aiohttp session if open.""" if self.session and not self.session.closed: await self.session.close() self.session = None def _get_session(self) -> aiohttp.ClientSession: """Get or create aiohttp session. Returns: Active aiohttp session """ # If no session, create one if self.session is None: timeout = aiohttp.ClientTimeout(total=self.timeout) self.session = aiohttp.ClientSession(timeout=timeout) return self.session # If session exists, check if it's closed (handle real sessions only) # Mock sessions from tests won't have a boolean closed attribute try: if hasattr(self.session, 'closed') and self.session.closed is True: timeout = aiohttp.ClientTimeout(total=self.timeout) self.session = aiohttp.ClientSession(timeout=timeout) except (AttributeError, TypeError): # Mock session or unusual object, just use it as-is pass return self.session async def download_image( self, url: str, local_path: Path, skip_existing: bool = True, validate: bool = True ) -> bool: """Download an image from URL to local path. Args: url: Image URL local_path: Local file path to save image skip_existing: Skip download if file already exists validate: Validate image after download Returns: True if download successful, False otherwise Raises: ImageDownloadError: If download fails after retries """ # Check if file already exists if skip_existing and local_path.exists(): if local_path.stat().st_size >= self.min_file_size: logger.debug("Image already exists: %s", local_path) return True # Ensure parent directory exists local_path.parent.mkdir(parents=True, exist_ok=True) delay = self.retry_delay last_error = None for attempt in range(self.max_retries): try: logger.debug( "Downloading image from %s (attempt %d)", url, attempt + 1, ) # Use persistent session session = self._get_session() async with session.get(url) as resp: if resp.status == 404: logger.warning("Image not found: %s", url) return False resp.raise_for_status() # Download image data data = await resp.read() # Check file size if len(data) < self.min_file_size: raise ImageDownloadError( f"Downloaded file too small: {len(data)} bytes" ) # Write to file with open(local_path, "wb") as f: f.write(data) # Validate image if requested if validate and not self.validate_image(local_path): local_path.unlink(missing_ok=True) raise ImageDownloadError("Image validation failed") logger.info("Downloaded image to %s", local_path) return True except (aiohttp.ClientError, IOError, ImageDownloadError) as e: last_error = e if attempt < self.max_retries - 1: logger.warning( "Download failed (attempt %d): %s, retrying in %s", attempt + 1, e, delay, ) await asyncio.sleep(delay) delay *= 2 else: logger.error( "Download failed after %d attempts: %s", self.max_retries, e, ) raise ImageDownloadError( f"Failed to download image after {self.max_retries} attempts: {last_error}" ) async def download_poster( self, url: str, series_folder: Path, filename: str = "poster.jpg", skip_existing: bool = True ) -> bool: """Download poster image. Args: url: Poster URL series_folder: Series folder path filename: Output filename (default: poster.jpg) skip_existing: Skip if file exists Returns: True if successful """ local_path = series_folder / filename try: return await self.download_image(url, local_path, skip_existing) except ImageDownloadError as e: logger.warning("Failed to download poster: %s", e) return False async def download_logo( self, url: str, series_folder: Path, filename: str = "logo.png", skip_existing: bool = True ) -> bool: """Download logo image. Args: url: Logo URL series_folder: Series folder path filename: Output filename (default: logo.png) skip_existing: Skip if file exists Returns: True if successful """ local_path = series_folder / filename try: return await self.download_image(url, local_path, skip_existing) except ImageDownloadError as e: logger.warning("Failed to download logo: %s", e) return False async def download_fanart( self, url: str, series_folder: Path, filename: str = "fanart.jpg", skip_existing: bool = True ) -> bool: """Download fanart/backdrop image. Args: url: Fanart URL series_folder: Series folder path filename: Output filename (default: fanart.jpg) skip_existing: Skip if file exists Returns: True if successful """ local_path = series_folder / filename try: return await self.download_image(url, local_path, skip_existing) except ImageDownloadError as e: logger.warning("Failed to download fanart: %s", e) return False def validate_image(self, image_path: Path) -> bool: """Validate that file is a valid image. Args: image_path: Path to image file Returns: True if valid image, False otherwise """ try: with Image.open(image_path) as img: # Verify it's a valid image img.verify() # Check file size if image_path.stat().st_size < self.min_file_size: logger.warning("Image file too small: %s", image_path) return False return True except Exception as e: logger.warning("Image validation failed for %s: %s", image_path, e) return False async def download_all_media( self, series_folder: Path, poster_url: Optional[str] = None, logo_url: Optional[str] = None, fanart_url: Optional[str] = None, skip_existing: bool = True ) -> dict[str, bool]: """Download all media files (poster, logo, fanart). Args: series_folder: Series folder path poster_url: Poster URL (optional) logo_url: Logo URL (optional) fanart_url: Fanart URL (optional) skip_existing: Skip existing files Returns: Dictionary with download status for each file type """ results = { "poster": None, "logo": None, "fanart": None } tasks = [] if poster_url: tasks.append(("poster", self.download_poster( poster_url, series_folder, skip_existing=skip_existing ))) if logo_url: tasks.append(("logo", self.download_logo( logo_url, series_folder, skip_existing=skip_existing ))) if fanart_url: tasks.append(("fanart", self.download_fanart( fanart_url, series_folder, skip_existing=skip_existing ))) # Download concurrently if tasks: task_results = await asyncio.gather( *[task for _, task in tasks], return_exceptions=True ) for (media_type, _), result in zip(tasks, task_results): if isinstance(result, Exception): logger.error("Error downloading %s: %s", media_type, result) results[media_type] = False else: results[media_type] = result return results