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 from abc import ABC, abstractmethod
   8 import os
   9 import os.path
  10
  11 from fetch import Fetcher
  12
  13
  14 class ImageStore(ABC):
  15
  16     @abstractmethod
  17     def get_image(self, url: str) -> str:
  18         raise NotImplementedError()
  19
  20
  21 class ImageStoreError(Exception):
  22     pass
  23
  24
  25 class DiskImageStore(ImageStore):
  26
  27     def __init__(self, root_path: str, fetcher: Fetcher) -> None:
  28         self._root_path = root_path
  29         self._fetcher = fetcher
  30         self._images: dict[str, str] = {}
  31         self._filenames: set[str] = set()
  32
  33     def _filename(self, url: str) -> str:
  34         assert url not in self._images
  35         base = os.path.basename(url).replace('%', '').split('?')[0]
  36         if base not in self._filenames:
  37             return base
  38         stem, ext = os.path.splitext(base)
  39         for i in range(10000):
  40             name = f"{stem}-{i:04d}{ext}"
  41             if name not in self._filenames:
  42                 return name
  43         raise ImageStoreError(
  44             'Unexpectedly-many similarly-named images fetched?')
  45
  46     def get_image(self, url: str) -> str:
  47         if url not in self._images:
  48             image_data = self._fetcher.fetch(url)
  49             name = self._filename(url)
  50             path = os.path.join(self._root_path, name)
  51             os.makedirs(self._root_path, exist_ok=True)
  52             with open(path, "wb") as f:
  53                 f.write(image_data)
  54             self._filenames.add(name)
  55             self._images[url] = path
  56         return self._images[url]
  57
  58
  59 class FakeImageStore(ImageStore):
  60
  61     def get_image(self, url: str) -> str:
  62         return 'stored:' + url