From 0e5d019a8a925d09a5d1ce3bb245b51f22a4c439 Mon Sep 17 00:00:00 2001 From: Kentai Radiquum Date: Mon, 20 Jan 2025 06:25:16 +0500 Subject: [PATCH] feat: add progress bar to the stdout for import process format logging close connections with telegram and matrix --- requirements.txt | 3 +- stickerbridge/matrix_reuploader.py | 46 +++++++++++++++--------------- stickerbridge/sticker-cli.py | 7 +++-- stickerbridge/sticker_types.py | 1 - stickerbridge/telegram_exporter.py | 24 ++++++++++------ 5 files changed, 46 insertions(+), 35 deletions(-) diff --git a/requirements.txt b/requirements.txt index cc607c4..6af27b6 100644 --- a/requirements.txt +++ b/requirements.txt @@ -5,4 +5,5 @@ pyyaml aiofiles python-magic lottie -cairosvg \ No newline at end of file +cairosvg +tqdm \ No newline at end of file diff --git a/stickerbridge/matrix_reuploader.py b/stickerbridge/matrix_reuploader.py index 06b4a36..cf951fa 100644 --- a/stickerbridge/matrix_reuploader.py +++ b/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 diff --git a/stickerbridge/sticker-cli.py b/stickerbridge/sticker-cli.py index ce0ed32..8266671 100644 --- a/stickerbridge/sticker-cli.py +++ b/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): diff --git a/stickerbridge/sticker_types.py b/stickerbridge/sticker_types.py index c385a69..25ea83e 100644 --- a/stickerbridge/sticker_types.py +++ b/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"], diff --git a/stickerbridge/telegram_exporter.py b/stickerbridge/telegram_exporter.py index 1b062e5..415ed8d 100644 --- a/stickerbridge/telegram_exporter.py +++ b/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