1 files 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: