Add a conversation handler for custom text color

overlay.py

@@ -3,7 +3,7 @@ from PIL import Image, ImageDraw, ImageFont
 def sub(p1, p2):
     return (p1[0] - p2[0], p1[1] - p2[1])
 
-def overlay(image_in, message="YOU DIED"):
+def overlay(image_in, message="YOU DIED", color=(200, 25, 25)):
     image = image_in.convert("RGBA")
 
     size = image.size
@@ -43,7 +43,7 @@ def overlay(image_in, message="YOU DIED"):
 
     offset = (center[0] - text_size[0]/2, center[1] - text_size[1]/2)
 
-    draw.text(offset, message, (200,25,25), font=font)
+    draw.text(offset, message, color, font=font)
 
     result = Image.alpha_composite(Image.alpha_composite(image, Image.new("RGBA", size, (0,0,0,125))), overlay)
     return result

you-died.py

@@ -1,5 +1,12 @@
-from telegram.ext import Updater, MessageHandler, Filters, CallbackContext
-from telegram import Update,Message
+from telegram.ext import (
+    Updater,
+    MessageHandler,
+    Filters,
+    CallbackContext,
+    ConversationHandler,
+    CommandHandler,
+)
+from telegram import Update, Message, ReplyKeyboardMarkup, ReplyKeyboardRemove, File
 import json
 import sys
 import io
@@ -8,8 +15,10 @@ from PIL import Image
 
 from overlay import overlay
 
-logging.basicConfig(format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
-                     level=logging.INFO)
+logging.basicConfig(
+    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", level=logging.INFO
+)
+
 
 def overlay_photo(update: Update, context: CallbackContext):
     try:
@@ -25,7 +34,11 @@ def overlay_photo(update: Update, context: CallbackContext):
         output.seek(0)
         context.bot.send_photo(chat_id=update.effective_chat.id, photo=output)
     except:
-        context.bot.send_message(chat_id=update.effective_chat.id, text="Something went wrong. That might not have been an image.")
+        context.bot.send_message(
+            chat_id=update.effective_chat.id,
+            text="Something went wrong. That might not have been an image.",
+        )
+
 
 def overlay_document(update: Update, context: CallbackContext):
     try:
@@ -41,10 +54,102 @@ def overlay_document(update: Update, context: CallbackContext):
         output.seek(0)
         context.bot.send_document(chat_id=update.effective_chat.id, document=output)
     except:
-        context.bot.send_message(chat_id=update.effective_chat.id, text="Something went wrong. That might not have been an image.")
+        context.bot.send_message(
+            chat_id=update.effective_chat.id,
+            text="Something went wrong. That might not have been an image.",
+        )
+
 
 def echo(update: Update, context: CallbackContext):
-    context.bot.send_message(chat_id=update.effective_chat.id, text="Send an image (inline or as a file). You can customize the text by adding a caption.")
+    context.bot.send_message(
+        chat_id=update.effective_chat.id,
+        text="Send an image (inline or as a file). You can customize the text by adding a caption. Alternatively, use the /caption command to choose the image, text, and color.",
+    )
+
+
+IMAGE, TEXT, COLOR = range(3)
+
+
+def start_overlay(update: Update, context: CallbackContext) -> int:
+    """Begins the overlay process."""
+    update.message.reply_text("First, send the image." "\n\n" "Send /cancel to stop.")
+
+    return IMAGE
+
+
+def receive_image(update: Update, context: CallbackContext) -> int:
+    """Receives the image to add text to."""
+    user = update.message.from_user
+    image = update.message.photo[-1].get_file()
+    context.user_data["image"] = image
+    logging.info("Received photo from %s", user.first_name)
+
+    update.message.reply_text("Next, send the caption.")
+
+    return TEXT
+
+
+def receive_text(update: Update, context: CallbackContext) -> int:
+    """Receives the caption to put on the image."""
+    user = update.message.from_user
+    text = update.message.text
+    context.user_data["text"] = text
+    options = [["Red", "Gold"]]
+
+    logging.info("Received text from %s", user.first_name)
+
+    update.message.reply_text(
+        "Finally, what color do you want?",
+        reply_markup=ReplyKeyboardMarkup(
+            options, one_time_keyboard=True, input_field_placeholder="Color"
+        ),
+    )
+
+    return COLOR
+
+
+def receive_color(update: Update, context: CallbackContext) -> int:
+    """Receives the color for the caption."""
+    user = update.message.from_user
+    color = update.message.text
+
+    logging.info("Received color from %s", user.first_name)
+
+    image: File = context.user_data["image"]
+    text: str = context.user_data["text"]
+
+    try:
+        bytes = image.download_as_bytearray()
+        image = Image.open(io.BytesIO(bytes))
+        color_rgb = (200, 25, 25)
+
+        if color == "Gold":
+            color_rgb = (255, 255, 108)
+
+        overlayed = overlay(image, text, color_rgb).convert("RGB")
+        output = io.BytesIO()
+        overlayed.save(output, "JPEG")
+        output.seek(0)
+        context.bot.send_photo(chat_id=update.effective_chat.id, photo=output)
+    except:
+        context.bot.send_message(
+            chat_id=update.effective_chat.id,
+            text="Something went wrong. That might not have been an image.",
+        )
+
+    update.message.reply_text("Done!", reply_markup=ReplyKeyboardRemove())
+
+    return ConversationHandler.END
+
+
+def cancel_overlay(update: Update, context: CallbackContext) -> int:
+    """Cancels the overlay process."""
+    user = update.message.from_user
+    logging.info("Canceled overlay for %s", user.first_name)
+    update.message.reply_text("Cancelled.", reply_markup=ReplyKeyboardRemove())
+
+    return ConversationHandler.END
+
 
 if __name__ == "__main__":
     try:
@@ -59,8 +164,20 @@ if __name__ == "__main__":
     overlay_document_handler = MessageHandler(Filters.document, overlay_document)
     echo_handler = MessageHandler(Filters.text, echo)
 
+    conv_handler = ConversationHandler(
+        entry_points=[CommandHandler("caption", start_overlay)],
+        states={
+            IMAGE: [MessageHandler(Filters.photo, receive_image)],
+            TEXT: [MessageHandler(Filters.text, receive_text)],
+            COLOR: [MessageHandler(Filters.regex("^(Red|Gold)$"), receive_color)],
+        },
+        fallbacks=[CommandHandler("cancel", cancel_overlay)],
+    )
+
+    dispatcher.add_handler(conv_handler)
+
     dispatcher.add_handler(overlay_photo_handler)
     dispatcher.add_handler(overlay_document_handler)
     dispatcher.add_handler(echo_handler)
 
-    updater.start_polling()
+    updater.start_polling()