]>
Commit | Line | Data |
---|---|---|
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 | ||
4b7d905e | 7 | from abc import ABC, abstractmethod |
91fe9916 SW |
8 | import os |
9 | import os.path | |
10 | ||
11 | from fetch import Fetcher | |
12 | ||
13 | ||
4b7d905e SW |
14 | class ImageStore(ABC): |
15 | ||
16 | @abstractmethod | |
17 | def get_image(self, url: str) -> str: | |
18 | raise NotImplementedError() | |
19 | ||
20 | ||
91fe9916 SW |
21 | class ImageStoreError(Exception): |
22 | pass | |
23 | ||
24 | ||
4b7d905e | 25 | class DiskImageStore(ImageStore): |
91fe9916 SW |
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 | |
a8a47f78 | 35 | base = os.path.basename(url).replace('%', '').split('?')[0] |
91fe9916 SW |
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] | |
4b7d905e SW |
57 | |
58 | ||
59 | class FakeImageStore(ImageStore): | |
60 | ||
61 | def get_image(self, url: str) -> str: | |
62 | return 'stored:' + url |