git.scottworley.com Git - paperdoorknob/blob - images.py

   1 # paperdoorknob: Print glowfic
   2 #
   3 # This program is free software: you can redistribute it and/or modify it
   4 # under the terms of the GNU General Public License as published by the
   5 # Free Software Foundation, version 3.
   6
   7 import os
   8 import os.path
   9
  10 from fetch import Fetcher
  11
  12
  13 class ImageStoreError(Exception):
  14     pass
  15
  16
  17 class ImageStore:
  18
  19     def __init__(self, root_path: str, fetcher: Fetcher) -> None:
  20         self._root_path = root_path
  21         self._fetcher = fetcher
  22         self._images: dict[str, str] = {}
  23         self._filenames: set[str] = set()
  24
  25     def _filename(self, url: str) -> str:
  26         assert url not in self._images
  27         base = os.path.basename(url)
  28         if base not in self._filenames:
  29             return base
  30         stem, ext = os.path.splitext(base)
  31         for i in range(10000):
  32             name = f"{stem}-{i:04d}{ext}"
  33             if name not in self._filenames:
  34                 return name
  35         raise ImageStoreError(
  36             'Unexpectedly-many similarly-named images fetched?')
  37
  38     def get_image(self, url: str) -> str:
  39         if url not in self._images:
  40             image_data = self._fetcher.fetch(url)
  41             name = self._filename(url)
  42             path = os.path.join(self._root_path, name)
  43             os.makedirs(self._root_path, exist_ok=True)
  44             with open(path, "wb") as f:
  45                 f.write(image_data)
  46             self._filenames.add(name)
  47             self._images[url] = path
  48         return self._images[url]