]> git.scottworley.com Git - paperdoorknob/blame_incremental - images.py
Use a more URL-looking url in tests
[paperdoorknob] / images.py
... / ...
CommitLineData
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
7from abc import ABC, abstractmethod
8import os
9import os.path
10
11from fetch import Fetcher
12
13
14class ImageStore(ABC):
15
16 @abstractmethod
17 def get_image(self, url: str) -> str:
18 raise NotImplementedError()
19
20
21class ImageStoreError(Exception):
22 pass
23
24
25class 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
59class FakeImageStore(ImageStore):
60
61 def get_image(self, url: str) -> str:
62 return 'stored:' + url