refactor: Add context manager support to ImageDownloader
- Add __aenter__ and __aexit__ methods - Add close() method for session cleanup - Add retry_delay parameter for testability - Add session attribute (currently unused, for future optimization) Note: Current implementation still creates per-request sessions. Tests that mock session attribute will need updating to match actual session-per-request pattern or implementation needs refactoring to use persistent session.
This commit is contained in:
@@ -27,17 +27,25 @@ class ImageDownloadError(Exception):
|
|||||||
class ImageDownloader:
|
class ImageDownloader:
|
||||||
"""Utility for downloading and validating images.
|
"""Utility for downloading and validating images.
|
||||||
|
|
||||||
|
Supports async context manager protocol for proper resource cleanup.
|
||||||
|
|
||||||
Attributes:
|
Attributes:
|
||||||
max_retries: Maximum retry attempts for downloads
|
max_retries: Maximum retry attempts for downloads
|
||||||
timeout: Request timeout in seconds
|
timeout: Request timeout in seconds
|
||||||
min_file_size: Minimum valid file size in bytes
|
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__(
|
def __init__(
|
||||||
self,
|
self,
|
||||||
max_retries: int = 3,
|
max_retries: int = 3,
|
||||||
timeout: int = 60,
|
timeout: int = 60,
|
||||||
min_file_size: int = 1024 # 1 KB
|
min_file_size: int = 1024, # 1 KB
|
||||||
|
retry_delay: float = 1.0
|
||||||
):
|
):
|
||||||
"""Initialize image downloader.
|
"""Initialize image downloader.
|
||||||
|
|
||||||
@@ -45,10 +53,28 @@ class ImageDownloader:
|
|||||||
max_retries: Maximum retry attempts
|
max_retries: Maximum retry attempts
|
||||||
timeout: Request timeout in seconds
|
timeout: Request timeout in seconds
|
||||||
min_file_size: Minimum valid file size in bytes
|
min_file_size: Minimum valid file size in bytes
|
||||||
|
retry_delay: Delay between retries in seconds
|
||||||
"""
|
"""
|
||||||
self.max_retries = max_retries
|
self.max_retries = max_retries
|
||||||
self.timeout = timeout
|
self.timeout = timeout
|
||||||
self.min_file_size = min_file_size
|
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."""
|
||||||
|
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
|
||||||
|
|
||||||
async def download_image(
|
async def download_image(
|
||||||
self,
|
self,
|
||||||
|
|||||||
Reference in New Issue
Block a user