56 lines
1.9 KiB
Python
56 lines
1.9 KiB
Python
"""Snapshot file storage and retention management."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class SnapshotStore:
|
|
"""Manages screenshot storage and cleanup."""
|
|
|
|
def __init__(self, base_directory: str, retention_count: int = 50) -> None:
|
|
self.base_directory = Path(base_directory)
|
|
self.retention_count = retention_count
|
|
self.base_directory.mkdir(parents=True, exist_ok=True)
|
|
|
|
def _website_directory(self, website_id: str) -> Path:
|
|
directory = self.base_directory / website_id
|
|
directory.mkdir(parents=True, exist_ok=True)
|
|
return directory
|
|
|
|
def _timestamp_filename(self, timestamp: datetime) -> str:
|
|
return timestamp.strftime("%Y-%m-%dT%H-%M-%S") + ".png"
|
|
|
|
def save_screenshot(
|
|
self,
|
|
website_id: str,
|
|
timestamp: datetime,
|
|
image_bytes: bytes,
|
|
) -> str:
|
|
"""Save a screenshot and enforce retention limits."""
|
|
directory = self._website_directory(website_id)
|
|
filename = self._timestamp_filename(timestamp)
|
|
filepath = directory / filename
|
|
filepath.write_bytes(image_bytes)
|
|
logger.debug("Screenshot saved: %s", filepath)
|
|
|
|
self._cleanup_old_snapshots(directory)
|
|
return str(filepath)
|
|
|
|
def _cleanup_old_snapshots(self, directory: Path) -> None:
|
|
"""Remove oldest snapshots beyond retention_count."""
|
|
screenshots = sorted(directory.glob("*.png"), key=lambda p: p.stat().st_mtime)
|
|
excess = len(screenshots) - self.retention_count
|
|
if excess <= 0:
|
|
return
|
|
|
|
for old_file in screenshots[:excess]:
|
|
try:
|
|
old_file.unlink()
|
|
logger.debug("Removed old snapshot: %s", old_file)
|
|
except OSError as exc:
|
|
logger.warning("Failed to remove snapshot %s: %s", old_file, exc)
|