Store IDs as strings and keep all emoticons. Fixes #12

This commit is contained in:
Tulir Asokan 2020-09-06 16:59:25 +03:00
parent c106c99b7f
commit 522e45a4ae

View File

@ -18,7 +18,7 @@ from PIL import Image
from telethon import TelegramClient from telethon import TelegramClient
from telethon.tl.functions.messages import GetAllStickersRequest, GetStickerSetRequest from telethon.tl.functions.messages import GetAllStickersRequest, GetStickerSetRequest
from telethon.tl.types.messages import AllStickers from telethon.tl.types.messages import AllStickers
from telethon.tl.types import InputStickerSetShortName, Document from telethon.tl.types import InputStickerSetShortName, Document, DocumentAttributeSticker
from telethon.tl.types.messages import StickerSet as StickerSetFull from telethon.tl.types.messages import StickerSet as StickerSetFull
parser = argparse.ArgumentParser() parser = argparse.ArgumentParser()
@ -71,6 +71,9 @@ async def upload(data: bytes, mimetype: str, filename: str) -> str:
if TYPE_CHECKING: if TYPE_CHECKING:
from typing import TypedDict
class MatrixMediaInfo(TypedDict): class MatrixMediaInfo(TypedDict):
w: int w: int
h: int h: int
@ -110,8 +113,12 @@ async def reupload_document(client: TelegramClient, document: Document) -> 'Matr
else: else:
width = int(width / (height / 256)) width = int(width / (height / 256))
height = 256 height = 256
body = ""
for attr in document.attributes:
if isinstance(attr, DocumentAttributeSticker):
body = attr.alt
return { return {
"body": "", "body": body,
"url": mxc, "url": mxc,
"info": { "info": {
"w": width, "w": width,
@ -163,7 +170,7 @@ async def reupload_pack(client: TelegramClient, pack: StickerSetFull) -> None:
try: try:
with open(pack_path) as pack_file: with open(pack_path) as pack_file:
existing_pack = json.load(pack_file) existing_pack = json.load(pack_file)
already_uploaded = {sticker["net.maunium.telegram.sticker"]["id"]: sticker already_uploaded = {int(sticker["net.maunium.telegram.sticker"]["id"]): sticker
for sticker in existing_pack["stickers"]} for sticker in existing_pack["stickers"]}
print(f"Found {len(already_uploaded)} already reuploaded stickers") print(f"Found {len(already_uploaded)} already reuploaded stickers")
except FileNotFoundError: except FileNotFoundError:
@ -176,26 +183,26 @@ async def reupload_pack(client: TelegramClient, pack: StickerSetFull) -> None:
print(f"Skipped reuploading {document.id}") print(f"Skipped reuploading {document.id}")
except KeyError: except KeyError:
reuploaded_documents[document.id] = await reupload_document(client, document) reuploaded_documents[document.id] = await reupload_document(client, document)
reuploaded_documents[document.id]["net.maunium.telegram.sticker"] = {
"pack": {
"id": str(pack.set.id),
"short_name": pack.set.short_name,
},
"id": str(document.id),
"emoticons": [],
}
for sticker in pack.packs: for sticker in pack.packs:
for document_id in sticker.documents: for document_id in sticker.documents:
doc = reuploaded_documents[document_id] doc = reuploaded_documents[document_id]
doc["body"] = sticker.emoticon doc["net.maunium.telegram.sticker"]["emoticons"].append(sticker.emoticon)
doc["net.maunium.telegram.sticker"] = {
"pack": {
"id": pack.set.id,
"short_name": pack.set.short_name,
},
"id": document_id,
"emoticon": sticker.emoticon,
}
with open(pack_path, "w") as pack_file: with open(pack_path, "w") as pack_file:
json.dump({ json.dump({
"title": pack.set.title, "title": pack.set.title,
"short_name": pack.set.short_name, "short_name": pack.set.short_name,
"id": pack.set.id, "id": str(pack.set.id),
"hash": pack.set.hash, "hash": str(pack.set.hash),
"stickers": list(reuploaded_documents.values()), "stickers": list(reuploaded_documents.values()),
}, pack_file, ensure_ascii=False) }, pack_file, ensure_ascii=False)
print(f"Saved {pack.set.title} as {pack.set.short_name}.json") print(f"Saved {pack.set.title} as {pack.set.short_name}.json")