]> git.scottworley.com Git - paperdoorknob/blob - images.py
178f7085f5a10a4feeae780b04707aacc92b07ae
[paperdoorknob] / 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]