feat: add progress bar to the stdout for import process

format logging
close connections with telegram and matrix

requirements.txt

@@ -5,4 +5,5 @@ pyyaml
 aiofiles
 python-magic
 lottie
-cairosvg
+cairosvg
+tqdm

stickerbridge/matrix_reuploader.py

@@ -3,6 +3,8 @@ import os
 import json
 import yaml
 import hashlib
+import logging
+from tqdm.auto import tqdm
 
 from nio import MatrixRoom, AsyncClient
 
@@ -57,7 +59,7 @@ async def _parse_args(args: list) -> dict[str, str]:
                 if not value.startswith("http"):
                     continue
                 parsed_args["artist_url"] = value
-        if arg in ["-p", "--primary", "-j", "--json"]:
+        if arg in ["-p", "--primary", "-j", "--json", "-upd", "--update-pack"]:
             if arg in ["-p", "--primary"]:
                 parsed_args["default"] = not parsed_args["default"]
             if arg in ["-j", "--json"]:
@@ -128,31 +130,29 @@ class MatrixReuploader:
         stickerset = MatrixStickerset(import_name, pack_name, parsed_args["rating"], {"name": parsed_args["artist"], "url": parsed_args["artist_url"]})
         json_stickerset = MauniumStickerset(import_name, pack_name, parsed_args["rating"], {"name": parsed_args["artist"], "url": parsed_args["artist_url"]}, self.room.room_id)
 
-        n = 0
-        for sticker in converted_stickerset:
-            with tempfile.NamedTemporaryFile('w+b', delete=False) as file:
-                file.write(sticker.image_data)
-                hash = hashlib.md5(sticker.image_data).hexdigest()
-                name = f"{pack_name}__{sticker.alt_text}__{os.path.basename(file.name)}"
+        with tqdm(total=len(converted_stickerset)) as tqdm_object:
+            for sticker in converted_stickerset:
+                with tempfile.NamedTemporaryFile('w+b', delete=False) as file:
+                    file.write(sticker.image_data)
+                    hash = hashlib.md5(sticker.image_data).hexdigest()
+                    name = f"{pack_name}__{sticker.alt_text}__{os.path.basename(file.name)}"
 
-                sticker_mxc = None
-                if stickerpack is not None and stickerpack.get('images', None) is not None:
-                    for stick in stickerpack['images'].values():
-                        if stick.get('hash', None) is not None and stick["hash"] == hash:
-                            sticker_mxc = stick["url"]
-                            print(f"sticker already exists, hash: {hash}")
-                            break
+                    sticker_mxc = None
+                    if stickerpack is not None and stickerpack.get('images', None) is not None:
+                        for stick in stickerpack['images'].values():
+                            if stick.get('hash', None) is not None and stick["hash"] == hash:
+                                sticker_mxc = stick["url"]
+                                break
 
-                if sticker_mxc is None:
-                    sticker_mxc = await upload_image(self.client, file.name, name)
-                file.close()
-                os.unlink(file.name)
+                    if sticker_mxc is None:
+                        sticker_mxc = await upload_image(self.client, file.name, name)
+                    file.close()
+                    os.unlink(file.name)
 
-            stickerset.add_sticker(sticker_mxc, sticker.alt_text, hash)
-            n += 1
-            print(f"Uploaded: {n}/{len(converted_stickerset)}")
-            if parsed_args["json"]:
-                json_stickerset.add_sticker(sticker_mxc, sticker.alt_text, sticker.width, sticker.height, sticker.size, sticker.mimetype)
+                stickerset.add_sticker(sticker_mxc, sticker.alt_text, hash)
+                if parsed_args["json"]:
+                    json_stickerset.add_sticker(sticker_mxc, sticker.alt_text, sticker.width, sticker.height, sticker.size, sticker.mimetype)
+                tqdm_object.update(1)
 
         if not stickerset.count():
             yield self.STATUS_PACK_EMPTY

stickerbridge/sticker-cli.py

@@ -39,7 +39,7 @@ import_cmd.add_argument('--rating', '-r', choices=('S', 'Q', 'E', 'U'), help='Se
 import_cmd.add_argument('--room', '-rm', type=str, help='Set a room for the sticker upload')
 import_cmd.add_argument('--create-room', '-cr', action='store_true', help='Create a new room for imported stickers')
 import_cmd.add_argument('--space', '-s', type=str, help='Space to include the new room in. (You will need to invite the bot first!)')
-import_cmd.add_argument('--update-pack', '-u', action='store_true', help='Update pack if it already exists')
+import_cmd.add_argument('--update-pack', '-upd', action='store_true', help='Update pack if it already exists')
 
 import_cmd.epilog = 'IF boolean flags are true in "config.yaml" or "cli.yaml", and are provided here, they are applied as a False.'
 
@@ -57,7 +57,8 @@ async def main(args):
     with open(args.cli_config, 'r') as config_file:
         cli_config = yaml.safe_load(config_file)
 
-    logging.basicConfig(level=os.environ.get("LOGLEVEL", config['log_level']))
+    fmt = f"%(asctime)-20s | %(filename)-20s | %(levelname)s : %(message)s"
+    logging.basicConfig(level=os.environ.get("LOGLEVEL", config['log_level']), format=fmt, handlers=[logging.StreamHandler()])
 
     client = AsyncClient(config['matrix_homeserver'], config['matrix_username'])
     client.device_id = config['matrix_bot_name']
@@ -74,6 +75,7 @@ async def main(args):
     if sys.argv[1] == 'import':
         await import_stickerpack(args, client, config, cli_config)
 
+    await client.close()
 
 async def import_stickerpack(args: argparse.Namespace, client: AsyncClient, config: dict, cli_config: dict):
     if args.pack_name.startswith('https://t.me/addstickers/'):
@@ -144,6 +146,7 @@ async def import_stickerpack(args: argparse.Namespace, client: AsyncClient, conf
             text = switch.get(status, "Warning: Unknown status")
             logging.info(text)
 
+    await tg_exporter.close()
 
 async def create_or_get_room(args: argparse.Namespace, client: AsyncClient, config: dict, cli_config: dict):
 

stickerbridge/sticker_types.py

@@ -29,7 +29,6 @@ class MatrixStickerset:
             while (alt_text in self._content['images']):
                 duplicate_counter += 1
                 alt_text = alt_text.split('-')[0] + '-' + str(duplicate_counter)
-        print(alt_text)
         self._content['images'][alt_text] = {
             "url": mxc_uri,
             "usage": ["sticker"],

stickerbridge/telegram_exporter.py

@@ -1,6 +1,9 @@
 from multiprocessing import Pool
 from typing import List
 
+import logging
+from tqdm.auto import tqdm
+
 from lottie.importers import importers
 from lottie.exporters import exporters
 from telethon import TelegramClient
@@ -71,7 +74,12 @@ class TelegramExporter:
     async def connect(self):
         await self.client.start(bot_token=self.bot_token)
 
+    async def close(self):
+        await self.client.disconnect()
+
     async def get_stickerset(self, pack_name: str) -> list[Sticker]:
+        logging.getLogger('telethon').setLevel(logging.WARNING)
+
         result: List[Sticker] = list()
 
         try:
@@ -81,15 +89,15 @@ class TelegramExporter:
 
         downloaded_documents = []
 
-        n = 0
-        for document_data in sticker_set.documents:
-            document_data.downloaded_data_ = await self.client.download_media(document_data, file=bytes)
-            downloaded_documents.append(document_data)
-            n += 1
-            print(f"Downloaded: {n}/{len(sticker_set.documents)}")
+        with tqdm(total=len(sticker_set.documents)) as tqdm_object:
+            for document_data in sticker_set.documents:
+                document_data.downloaded_data_ = await self.client.download_media(document_data, file=bytes)
+                downloaded_documents.append(document_data)
+                tqdm_object.update(1)
 
-        print("Processing stickers...")
+        logging.info(f"Processing downloaded stickers...")
         pool = Pool()
-        result = pool.map(_process_sticker, downloaded_documents)
+        # result = pool.map(_process_sticker, downloaded_documents)
+        result = list(tqdm(pool.imap(_process_sticker, downloaded_documents), total=len(downloaded_documents)))
 
         return result