]> git.scottworley.com Git - paperdoorknob/blame - images.py
ImageStore
[paperdoorknob] / images.py
CommitLineData
91fe9916
SW
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
7import os
8import os.path
9
10from fetch import Fetcher
11
12
13class ImageStoreError(Exception):
14 pass
15
16
17class 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]