first commit

storage/__init__.py

@@ -0,0 +1,5 @@
+"""Storage package."""
+
+from storage.snapshots import SnapshotStore
+
+__all__ = ["SnapshotStore"]

storage/snapshots.py

@@ -0,0 +1,56 @@
+"""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)