telegram_bots/bots/modules/overlay.py

102 lines
4.2 KiB
Python
Raw Normal View History

import os
from tempfile import TemporaryDirectory
from time import sleep
from PIL import Image
from django.db import models
2020-02-10 15:57:55 +00:00
from django import forms
from django.forms import modelformset_factory
2020-06-02 12:50:33 +00:00
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)
2020-06-02 12:50:33 +00:00
send_avatar_on_start = models.BooleanField(default=False)
MODULE_NAME = 'Overlay'
def start_handler(self, update: Update, ctx: CallbackContext):
2020-06-02 12:50:33 +00:00
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)
2020-06-02 12:50:33 +00:00
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
2020-02-10 15:57:55 +00:00
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
2020-02-10 15:57:55 +00:00
@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())
2020-02-10 16:07:49 +00:00
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'])
2020-02-10 15:57:55 +00:00
class OverlayImage(models.Model):
config = models.ForeignKey(OverlayBotModuleConfig, on_delete=models.CASCADE, related_name='images')
image = models.ImageField()
class ImageForm(forms.ModelForm):
2020-02-10 16:02:45 +00:00
image = forms.ImageField(widget=forms.ClearableFileInput(attrs={'accept': 'image/png'}), required=False)
2020-02-10 15:57:55 +00:00
class Meta:
model = OverlayImage
fields = 'image',