bot: process images in memory

1 files changed, 29 insertions, 23 deletions

diff --git a/bot.py b/bot.py
index 2d48783..bf19588 100644
--- a/bot.py
+++ b/bot.py
@@ -1,5 +1,7 @@
+import io
 import logging
 import os
+import traceback
 
 from telegram import __version__ as TG_VER
 
@@ -19,8 +21,8 @@ if __version_info__ < (20, 0, 0, "alpha", 1):
 from telegram import ForceReply, Update, File
 from telegram.ext import Application, CommandHandler, ContextTypes, MessageHandler, filters
 from telegram.constants import ParseMode
-from PIL import Image, ImageFilter
-from square import square_size_padding, drop_shadow
+from PIL import Image
+from square import square_size_padding
 
 # Enable logging
 
@@ -31,52 +33,56 @@ logging.basicConfig(
 logger = logging.getLogger(__name__)
 
 
-# Define a few command handlers. These usually take the two arguments update and context.
-
 async def start(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
-    """Send a message when the command /start is issued."""
-    user = update.effective_user
-    await update.message.reply_html(
-        rf"Hi {user.mention_html()}!",
-    )
+    await update.message.reply_text("This is a bot to output image in square shape")
 
 
 async def process(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
-    """Echo the user message."""
-    filename = update.message.document.file_name
     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])
+    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 = await file.download_to_drive(filename)
+    img = io.BytesIO()
+    await file.download_to_memory(img)
 
     try:
         im = Image.open(img)
-        names = filename.split(".")
         result = square_size_padding(im)
-        filename = "{}-result.{}".format(names[0], names[1])
-        result.save(filename, quality=100)
 
-        await context.bot.send_document(chat_id=update.message.chat_id, document=filename)
+        output = io.BytesIO()
+        result.save(output, format=save_format, quality=100)
+
+        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]),
+                                        document=output.getvalue())
 
     except Exception as e:
-        await update.message.reply_markdown_v2(text="Error:\n```{}```".format(str(e)))
+        await update.message.reply_markdown_v2(text="Error:\n```{}```".format(traceback.format_exc()))
 
 
 def main() -> None:
-    """Start the bot."""
-    # Create the Application and pass it your bot's token.
     tg_token = os.getenv("TG_TOKEN")
     application = Application.builder().token(tg_token).build()
 
-    # on different commands - answer in Telegram
     application.add_handler(CommandHandler("start", start))
-
-    # on non command i.e message - echo the message on Telegram
     application.add_handler(MessageHandler(filters.ATTACHMENT & ~filters.COMMAND, process))
 
-    # Run the bot until the user presses Ctrl-C
     application.run_polling()