mirror of
https://github.com/wah-su/stickerbridge.git
synced 2025-04-04 23:34:40 +00:00
feat: add progress bar to the stdout for import process
format logging close connections with telegram and matrix
This commit is contained in:
parent
b24dae6c18
commit
0e5d019a8a
5 changed files with 46 additions and 35 deletions
|
@ -5,4 +5,5 @@ pyyaml
|
||||||
aiofiles
|
aiofiles
|
||||||
python-magic
|
python-magic
|
||||||
lottie
|
lottie
|
||||||
cairosvg
|
cairosvg
|
||||||
|
tqdm
|
|
@ -3,6 +3,8 @@ import os
|
||||||
import json
|
import json
|
||||||
import yaml
|
import yaml
|
||||||
import hashlib
|
import hashlib
|
||||||
|
import logging
|
||||||
|
from tqdm.auto import tqdm
|
||||||
|
|
||||||
from nio import MatrixRoom, AsyncClient
|
from nio import MatrixRoom, AsyncClient
|
||||||
|
|
||||||
|
@ -57,7 +59,7 @@ async def _parse_args(args: list) -> dict[str, str]:
|
||||||
if not value.startswith("http"):
|
if not value.startswith("http"):
|
||||||
continue
|
continue
|
||||||
parsed_args["artist_url"] = value
|
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"]:
|
if arg in ["-p", "--primary"]:
|
||||||
parsed_args["default"] = not parsed_args["default"]
|
parsed_args["default"] = not parsed_args["default"]
|
||||||
if arg in ["-j", "--json"]:
|
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"]})
|
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)
|
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
|
with tqdm(total=len(converted_stickerset)) as tqdm_object:
|
||||||
for sticker in converted_stickerset:
|
for sticker in converted_stickerset:
|
||||||
with tempfile.NamedTemporaryFile('w+b', delete=False) as file:
|
with tempfile.NamedTemporaryFile('w+b', delete=False) as file:
|
||||||
file.write(sticker.image_data)
|
file.write(sticker.image_data)
|
||||||
hash = hashlib.md5(sticker.image_data).hexdigest()
|
hash = hashlib.md5(sticker.image_data).hexdigest()
|
||||||
name = f"{pack_name}__{sticker.alt_text}__{os.path.basename(file.name)}"
|
name = f"{pack_name}__{sticker.alt_text}__{os.path.basename(file.name)}"
|
||||||
|
|
||||||
sticker_mxc = None
|
sticker_mxc = None
|
||||||
if stickerpack is not None and stickerpack.get('images', None) is not None:
|
if stickerpack is not None and stickerpack.get('images', None) is not None:
|
||||||
for stick in stickerpack['images'].values():
|
for stick in stickerpack['images'].values():
|
||||||
if stick.get('hash', None) is not None and stick["hash"] == hash:
|
if stick.get('hash', None) is not None and stick["hash"] == hash:
|
||||||
sticker_mxc = stick["url"]
|
sticker_mxc = stick["url"]
|
||||||
print(f"sticker already exists, hash: {hash}")
|
break
|
||||||
break
|
|
||||||
|
|
||||||
if sticker_mxc is None:
|
if sticker_mxc is None:
|
||||||
sticker_mxc = await upload_image(self.client, file.name, name)
|
sticker_mxc = await upload_image(self.client, file.name, name)
|
||||||
file.close()
|
file.close()
|
||||||
os.unlink(file.name)
|
os.unlink(file.name)
|
||||||
|
|
||||||
stickerset.add_sticker(sticker_mxc, sticker.alt_text, hash)
|
stickerset.add_sticker(sticker_mxc, sticker.alt_text, hash)
|
||||||
n += 1
|
if parsed_args["json"]:
|
||||||
print(f"Uploaded: {n}/{len(converted_stickerset)}")
|
json_stickerset.add_sticker(sticker_mxc, sticker.alt_text, sticker.width, sticker.height, sticker.size, sticker.mimetype)
|
||||||
if parsed_args["json"]:
|
tqdm_object.update(1)
|
||||||
json_stickerset.add_sticker(sticker_mxc, sticker.alt_text, sticker.width, sticker.height, sticker.size, sticker.mimetype)
|
|
||||||
|
|
||||||
if not stickerset.count():
|
if not stickerset.count():
|
||||||
yield self.STATUS_PACK_EMPTY
|
yield self.STATUS_PACK_EMPTY
|
||||||
|
|
|
@ -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('--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('--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('--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.'
|
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:
|
with open(args.cli_config, 'r') as config_file:
|
||||||
cli_config = yaml.safe_load(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 = AsyncClient(config['matrix_homeserver'], config['matrix_username'])
|
||||||
client.device_id = config['matrix_bot_name']
|
client.device_id = config['matrix_bot_name']
|
||||||
|
@ -74,6 +75,7 @@ async def main(args):
|
||||||
if sys.argv[1] == 'import':
|
if sys.argv[1] == 'import':
|
||||||
await import_stickerpack(args, client, config, cli_config)
|
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):
|
async def import_stickerpack(args: argparse.Namespace, client: AsyncClient, config: dict, cli_config: dict):
|
||||||
if args.pack_name.startswith('https://t.me/addstickers/'):
|
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")
|
text = switch.get(status, "Warning: Unknown status")
|
||||||
logging.info(text)
|
logging.info(text)
|
||||||
|
|
||||||
|
await tg_exporter.close()
|
||||||
|
|
||||||
async def create_or_get_room(args: argparse.Namespace, client: AsyncClient, config: dict, cli_config: dict):
|
async def create_or_get_room(args: argparse.Namespace, client: AsyncClient, config: dict, cli_config: dict):
|
||||||
|
|
||||||
|
|
|
@ -29,7 +29,6 @@ class MatrixStickerset:
|
||||||
while (alt_text in self._content['images']):
|
while (alt_text in self._content['images']):
|
||||||
duplicate_counter += 1
|
duplicate_counter += 1
|
||||||
alt_text = alt_text.split('-')[0] + '-' + str(duplicate_counter)
|
alt_text = alt_text.split('-')[0] + '-' + str(duplicate_counter)
|
||||||
print(alt_text)
|
|
||||||
self._content['images'][alt_text] = {
|
self._content['images'][alt_text] = {
|
||||||
"url": mxc_uri,
|
"url": mxc_uri,
|
||||||
"usage": ["sticker"],
|
"usage": ["sticker"],
|
||||||
|
|
|
@ -1,6 +1,9 @@
|
||||||
from multiprocessing import Pool
|
from multiprocessing import Pool
|
||||||
from typing import List
|
from typing import List
|
||||||
|
|
||||||
|
import logging
|
||||||
|
from tqdm.auto import tqdm
|
||||||
|
|
||||||
from lottie.importers import importers
|
from lottie.importers import importers
|
||||||
from lottie.exporters import exporters
|
from lottie.exporters import exporters
|
||||||
from telethon import TelegramClient
|
from telethon import TelegramClient
|
||||||
|
@ -71,7 +74,12 @@ class TelegramExporter:
|
||||||
async def connect(self):
|
async def connect(self):
|
||||||
await self.client.start(bot_token=self.bot_token)
|
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]:
|
async def get_stickerset(self, pack_name: str) -> list[Sticker]:
|
||||||
|
logging.getLogger('telethon').setLevel(logging.WARNING)
|
||||||
|
|
||||||
result: List[Sticker] = list()
|
result: List[Sticker] = list()
|
||||||
|
|
||||||
try:
|
try:
|
||||||
|
@ -81,15 +89,15 @@ class TelegramExporter:
|
||||||
|
|
||||||
downloaded_documents = []
|
downloaded_documents = []
|
||||||
|
|
||||||
n = 0
|
with tqdm(total=len(sticker_set.documents)) as tqdm_object:
|
||||||
for document_data in sticker_set.documents:
|
for document_data in sticker_set.documents:
|
||||||
document_data.downloaded_data_ = await self.client.download_media(document_data, file=bytes)
|
document_data.downloaded_data_ = await self.client.download_media(document_data, file=bytes)
|
||||||
downloaded_documents.append(document_data)
|
downloaded_documents.append(document_data)
|
||||||
n += 1
|
tqdm_object.update(1)
|
||||||
print(f"Downloaded: {n}/{len(sticker_set.documents)}")
|
|
||||||
|
|
||||||
print("Processing stickers...")
|
logging.info(f"Processing downloaded stickers...")
|
||||||
pool = Pool()
|
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
|
return result
|
||||||
|
|
Loading…
Add table
Reference in a new issue