telegram_bots/bots/modules/overlay.py

import os
from tempfile import TemporaryDirectory
from time import sleep

from PIL import Image
from django.db import models
from django import forms
from django.forms import modelformset_factory
from telegram import Update, UserProfilePhotos
from telegram.ext import Dispatcher, CallbackContext, MessageHandler, Filters, CommandHandler

from bots.models import TelegramBotModuleConfig


class OverlayBotModuleConfig(TelegramBotModuleConfig):
    start_text = models.TextField(null=True, blank=True)
    type_error_text = models.TextField(null=True, blank=True)
    comment = models.TextField(blank=True, null=True)
    send_avatar_on_start = models.BooleanField(default=False)

    MODULE_NAME = 'Overlay'

    def start_handler(self, update: Update, ctx: CallbackContext):
        if self.send_avatar_on_start:
            with TemporaryDirectory() as d:
                src = os.path.join(d, 'src.jpg')
                photos = update.effective_user.get_profile_photos(limit=1)  # type: UserProfilePhotos
                if photos.photos:
                    q = photos.photos[0][-1].get_file().download(src)
                    im = Image.open(q)
                    self.create_image(update, im)
                elif self.type_error_text:
                    update.effective_message.reply_text(self.type_error_text)
        elif self.start_text:
            update.effective_message.reply_text(self.start_text)

    def message_handler(self, update: Update, ctx: CallbackContext):
        if self.type_error_text:
            update.effective_message.reply_text(self.type_error_text)

    def photo_handler(self, update: Update, ctx: CallbackContext):
        with TemporaryDirectory() as d:
            src = os.path.join(d, 'src.jpg')
            q = update.effective_message.photo[-1].get_file().download(src)
            im = Image.open(q)
            self.create_image(update, im)

    def create_image(self, update: Update, im: Image.Image):
        with TemporaryDirectory() as d:
            out = os.path.join(d, 'out.png')
            im = im.convert('RGBA')  # type: Image.Image
            overlay = Image.open(self.images.order_by('?').first().image.path).convert('RGBA')  # type: Image.Image
            w, h = im.size
            min_side = min(w, h)
            overlay = overlay.resize((min_side, min_side), Image.LANCZOS)
            im.paste(overlay, ((w - min_side) // 2, (h - min_side) // 2), overlay)
            im.save(out)
            update.effective_message.reply_photo(open(out, 'rb'), caption=self.comment)

    def build_dispatcher(self, dispatcher: Dispatcher):
        dispatcher.add_handler(CommandHandler('start', self.start_handler))
        dispatcher.add_handler(MessageHandler(Filters.photo, self.photo_handler))
        dispatcher.add_handler(MessageHandler(Filters.all, self.message_handler))
        return dispatcher

    @staticmethod
    def form_context_data(view, instance=None):
        ImageFormSet = modelformset_factory(OverlayImage, form=ImageForm, extra=3, can_delete=True)
        formset = ImageFormSet(queryset=instance.images.all() if instance else OverlayImage.objects.none())
        return {
            'formset': formset,
        }

    @staticmethod
    def form_valid(view, bots_form, config_form):
        ImageFormSet = modelformset_factory(OverlayImage, form=ImageForm, extra=3, can_delete=True)
        formset = ImageFormSet(view.request.POST, view.request.FILES, queryset=config_form.instance.images.all())
        for form in formset.cleaned_data:
            if form:
                if form['DELETE']:
                    form['id'].delete()
                elif form['image']:
                    if form['id']:
                        instance = form['id']
                        instance.image = form['image']
                        instance.save()
                    else:
                        OverlayImage.objects.create(config=config_form.instance, image=form['image'])


class OverlayImage(models.Model):
    config = models.ForeignKey(OverlayBotModuleConfig, on_delete=models.CASCADE, related_name='images')
    image = models.ImageField()


class ImageForm(forms.ModelForm):
    image = forms.ImageField(widget=forms.ClearableFileInput(attrs={'accept': 'image/png'}), required=False)

    class Meta:
        model = OverlayImage
        fields = 'image',