From 813834094d0cc8b476760be55e3db425416c0f63 Mon Sep 17 00:00:00 2001 From: clarkzjw Date: Fri, 25 Nov 2022 13:41:10 -0800 Subject: bot: support compressed photo --- bot.py | 37 +++++++++++++++++++++++-------------- 1 file changed, 23 insertions(+), 14 deletions(-) diff --git a/bot.py b/bot.py index 979b025..44505b1 100644 --- a/bot.py +++ b/bot.py @@ -10,8 +10,6 @@ from telegram.ext import Application, CommandHandler, ContextTypes, MessageHandl from square import square_size_padding -# Enable logging - logging.basicConfig( format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", level=logging.INFO ) @@ -25,23 +23,34 @@ async def start(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: async def process(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: chat_id = update.message.chat_id - - filename = update.message.document.file_name - names = filename.split(".") - - if str.upper(names[1]) in ("JPG", "JPEG"): - save_format = "JPEG" - elif str.upper(names[1]) in ("PNG",): - save_format = str.upper(names[1]) + save_format = "JPEG" + + if update.message.document is not None: + names = update.message.document.file_name.split(".") + file_ext = names[1] + filename = names[0] + + if str.upper(file_ext) in ("JPG", "JPEG"): + save_format = "JPEG" + elif str.upper(file_ext) in ("PNG",): + save_format = str.upper(file_ext) + else: + await context.bot.send_message(chat_id, "Image extension `{}` not supported".format(file_ext), + parse_mode=ParseMode.MARKDOWN_V2) + return + + file = await update.message.effective_attachment.get_file() + + elif update.message.photo is not None: + filename = update.message.photo[-1].file_unique_id + file_ext = "JPG" + file = await update.message.effective_attachment[-1].get_file() else: - await context.bot.send_message(chat_id, "Image extension `{}` not supported".format(names[1]), - parse_mode=ParseMode.MARKDOWN_V2) return await context.bot.send_message(chat_id, "Processing `{}`".format(filename), parse_mode=ParseMode.MARKDOWN_V2) - file = await update.message.effective_attachment.get_file() img = io.BytesIO() await file.download_to_memory(img) @@ -55,7 +64,7 @@ async def process(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: await update.message.reply_markdown_v2(text="Sending processed result") await context.bot.send_document(chat_id=update.message.chat_id, - filename="{}-result.{}".format(names[0], names[1]), + filename="{}-result.{}".format(filename, file_ext), document=output.getvalue()) except Exception as e: -- cgit v1.2.3