ImageStore

[paperdoorknob] / images.py

# paperdoorknob: Print glowfic
#
# This program is free software: you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by the
# Free Software Foundation, version 3.

import os
import os.path

from fetch import Fetcher


class ImageStoreError(Exception):
    pass


class ImageStore:

    def __init__(self, root_path: str, fetcher: Fetcher) -> None:
        self._root_path = root_path
        self._fetcher = fetcher
        self._images: dict[str, str] = {}
        self._filenames: set[str] = set()

    def _filename(self, url: str) -> str:
        assert url not in self._images
        base = os.path.basename(url)
        if base not in self._filenames:
            return base
        stem, ext = os.path.splitext(base)
        for i in range(10000):
            name = f"{stem}-{i:04d}{ext}"
            if name not in self._filenames:
                return name
        raise ImageStoreError(
            'Unexpectedly-many similarly-named images fetched?')

    def get_image(self, url: str) -> str:
        if url not in self._images:
            image_data = self._fetcher.fetch(url)
            name = self._filename(url)
            path = os.path.join(self._root_path, name)
            os.makedirs(self._root_path, exist_ok=True)
            with open(path, "wb") as f:
                f.write(image_data)
            self._filenames.add(name)
            self._images[url] = path
        return self._images[url]