ImageStore

[paperdoorknob] / images.py

diff --git a/images.py b/images.py
new file mode 100644 (file)
index 0000000..178f708
--- /dev/null
+++ b/images.py
@@ -0,0 +1,48 @@
+# paperdoorknob: Print glowfic
+#
+# This program is free software: you can redistribute it and/or modify it
+# under the terms of the GNU General Public License as published by the
+# Free Software Foundation, version 3.
+
+import os
+import os.path
+
+from fetch import Fetcher
+
+
+class ImageStoreError(Exception):
+    pass
+
+
+class ImageStore:
+
+    def __init__(self, root_path: str, fetcher: Fetcher) -> None:
+        self._root_path = root_path
+        self._fetcher = fetcher
+        self._images: dict[str, str] = {}
+        self._filenames: set[str] = set()
+
+    def _filename(self, url: str) -> str:
+        assert url not in self._images
+        base = os.path.basename(url)
+        if base not in self._filenames:
+            return base
+        stem, ext = os.path.splitext(base)
+        for i in range(10000):
+            name = f"{stem}-{i:04d}{ext}"
+            if name not in self._filenames:
+                return name
+        raise ImageStoreError(
+            'Unexpectedly-many similarly-named images fetched?')
+
+    def get_image(self, url: str) -> str:
+        if url not in self._images:
+            image_data = self._fetcher.fetch(url)
+            name = self._filename(url)
+            path = os.path.join(self._root_path, name)
+            os.makedirs(self._root_path, exist_ok=True)
+            with open(path, "wb") as f:
+                f.write(image_data)
+            self._filenames.add(name)
+            self._images[url] = path
+        return self._images[url]