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