utils/image_store.py

import json
import logging
from pathlib import Path

logger = logging.getLogger(__name__)
IMAGE_DIR = Path(__file__).parent.parent / "data" / "images"


class ImageStore:
    def __init__(self, image_dir=IMAGE_DIR):
        self.dir = Path(image_dir)
        self.dir.mkdir(parents=True, exist_ok=True)
        self._index_path = self.dir / "index.json"
        self._index = {}
        self._load_index()

    def _load_index(self):
        if self._index_path.exists():
            try:
                self._index = json.loads(self._index_path.read_text())
            except Exception:
                self._index = {}

    def _save_index(self):
        self._index_path.write_text(json.dumps(self._index, indent=2))

    def _source_dir(self, source: str) -> Path:
        safe = source.replace("/", "_").replace("\\", "_")
        return self.dir / safe

    def save(self, source: str, images: list):
        """Save list of (page, img_idx, PIL.Image) tuples. Returns count saved."""
        self._remove_files(source)
        self._index[source] = []
        src_dir = self._source_dir(source)
        src_dir.mkdir(parents=True, exist_ok=True)
        for page, img_idx, pil_img in images:
            path = src_dir / f"p{page}_i{img_idx}.png"
            pil_img.save(path, format="PNG")
            self._index[source].append({"path": str(path), "page": page, "image_index": img_idx})
        self._save_index()
        return len(images)

    def list_images(self, source: str) -> list:
        return [e for e in self._index.get(source, []) if Path(e["path"]).exists()]

    def has_images(self, source: str) -> bool:
        return bool(self._index.get(source))

    def was_attempted(self, source: str) -> bool:
        return source in self._index

    def remove(self, source: str):
        self._remove_files(source)
        self._index.pop(source, None)
        self._save_index()

    def clear_all(self):
        for source in list(self._index.keys()):
            self._remove_files(source)
        self._index = {}
        self._save_index()

    def _remove_files(self, source: str):
        for e in self._index.get(source, []):
            Path(e["path"]).unlink(missing_ok=True)
        src_dir = self._source_dir(source)
        if src_dir.exists() and not any(src_dir.iterdir()):
            src_dir.rmdir()