102 lines
4.2 KiB
Python
102 lines
4.2 KiB
Python
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',
|