Save and server thumbnails locally (#88)

Fixes #83
Fixes #86
Fixes #87

3 files changed, 32 insertions, 10 deletions

diff --git a/sticker/lib/util.py b/sticker/lib/util.py
index 2b3fe2a..1d2b84c 100644
--- a/sticker/lib/util.py
+++ b/sticker/lib/util.py
@@ -17,6 +17,8 @@ from functools import partial
 from io import BytesIO
 import os.path
 import json
+from pathlib import Path
+from typing import Dict, List
 
 from PIL import Image
 
@@ -24,19 +26,19 @@ from . import matrix
 
 open_utf8 = partial(open, encoding='UTF-8')
 
-def convert_image(data: bytes) -> (bytes, int, int):
+def convert_image(data: bytes, max_w=256, max_h=256) -> (bytes, int, int):
     image: Image.Image = Image.open(BytesIO(data)).convert("RGBA")
     new_file = BytesIO()
     image.save(new_file, "png")
     w, h = image.size
-    if w > 256 or h > 256:
+    if w > max_w or h > max_h:
         # Set the width and height to lower values so clients wouldn't show them as huge images
         if w > h:
-            h = int(h / (w / 256))
-            w = 256
+            h = int(h / (w / max_w))
+            w = max_w
         else:
-            w = int(w / (h / 256))
-            h = 256
+            w = int(w / (h / max_h))
+            h = max_h
     return new_file.getvalue(), w, h
 
 
@@ -78,3 +80,15 @@ def make_sticker(mxc: str, width: int, height: int, size: int,
         },
         "msgtype": "m.sticker",
     }
+
+
+def add_thumbnails(stickers: List[matrix.StickerInfo], stickers_data: Dict[str, bytes], output_dir: str) -> None:
+    thumbnails = Path(output_dir, "thumbnails")
+    thumbnails.mkdir(parents=True, exist_ok=True)
+
+    for sticker in stickers:       
+        image_data, _, _ = convert_image(stickers_data[sticker["url"]], 128, 128)
+        
+        name = sticker["url"].split("/")[-1]
+        thumbnail_path = thumbnails / name
+        thumbnail_path.write_bytes(image_data)
\ No newline at end of file
diff --git a/sticker/pack.py b/sticker/pack.py
index f082370..a4ca1f2 100644
--- a/sticker/pack.py
+++ b/sticker/pack.py
@@ -13,6 +13,7 @@
 #
 # You should have received a copy of the GNU Affero General Public License
 # along with this program.  If not, see <https://www.gnu.org/licenses/>.
+from pathlib import Path
 from typing import Dict, Optional
 from hashlib import sha256
 import mimetypes
@@ -107,9 +108,11 @@ async def main(args: argparse.Namespace) -> None:
         old_stickers = {sticker["id"]: sticker for sticker in pack["stickers"]}
         pack["stickers"] = []
 
+    stickers_data: Dict[str, bytes] = {}
     for file in sorted(os.listdir(args.path)):
         sticker = await upload_sticker(file, args.path, old_stickers=old_stickers)
         if sticker:
+            stickers_data[sticker["url"]]  = Path(args.path, file).read_bytes()
             pack["stickers"].append(sticker)
 
     with util.open_utf8(meta_path, "w") as pack_file:
@@ -122,6 +125,8 @@ async def main(args: argparse.Namespace) -> None:
         with util.open_utf8(picker_pack_path, "w") as pack_file:
             json.dump(pack, pack_file)
         print(f"Copied pack to {picker_pack_path}")
+
+        util.add_thumbnails(pack["stickers"], stickers_data, args.add_to_index)
         util.add_to_index(picker_file_name, args.add_to_index)
 
 
diff --git a/sticker/stickerimport.py b/sticker/stickerimport.py
index 534f3c4..54cab76 100644
--- a/sticker/stickerimport.py
+++ b/sticker/stickerimport.py
@@ -13,7 +13,7 @@
 #
 # You should have received a copy of the GNU Affero General Public License
 # along with this program.  If not, see <https://www.gnu.org/licenses/>.
-from typing import Dict
+from typing import Dict, Tuple
 import argparse
 import asyncio
 import os.path
@@ -29,7 +29,7 @@ from telethon.tl.types.messages import StickerSet as StickerSetFull
 from .lib import matrix, util
 
 
-async def reupload_document(client: TelegramClient, document: Document) -> matrix.StickerInfo:
+async def reupload_document(client: TelegramClient, document: Document) -> Tuple[matrix.StickerInfo, bytes]:
     print(f"Reuploading {document.id}", end="", flush=True)
     data = await client.download_media(document, file=bytes)
     print(".", end="", flush=True)
@@ -37,7 +37,7 @@ async def reupload_document(client: TelegramClient, document: Document) -> matri
     print(".", end="", flush=True)
     mxc = await matrix.upload(data, "image/png", f"{document.id}.png")
     print(".", flush=True)
-    return util.make_sticker(mxc, width, height, len(data))
+    return util.make_sticker(mxc, width, height, len(data)), data
 
 
 def add_meta(document: Document, info: matrix.StickerInfo, pack: StickerSetFull) -> None:
@@ -75,15 +75,17 @@ async def reupload_pack(client: TelegramClient, pack: StickerSetFull, output_dir
     except FileNotFoundError:
         pass
 
+    stickers_data: Dict[str, bytes] = {}
     reuploaded_documents: Dict[int, matrix.StickerInfo] = {}
     for document in pack.documents:
         try:
             reuploaded_documents[document.id] = already_uploaded[document.id]
             print(f"Skipped reuploading {document.id}")
         except KeyError:
-            reuploaded_documents[document.id] = await reupload_document(client, document)
+            reuploaded_documents[document.id], data = await reupload_document(client, document)
         # Always ensure the body and telegram metadata is correct
         add_meta(document, reuploaded_documents[document.id], pack)
+        stickers_data[reuploaded_documents[document.id]["url"]] = data
 
     for sticker in pack.packs:
         if not sticker.emoticon:
@@ -107,6 +109,7 @@ async def reupload_pack(client: TelegramClient, pack: StickerSetFull, output_dir
         }, pack_file, ensure_ascii=False)
     print(f"Saved {pack.set.title} as {pack.set.short_name}.json")
 
+    util.add_thumbnails(list(reuploaded_documents.values()), stickers_data, output_dir)
     util.add_to_index(os.path.basename(pack_path), output_dir)