From 752256be7bc37e27daa63b4d59524b57e72f3f99 Mon Sep 17 00:00:00 2001 From: bakatrouble Date: Sat, 16 Nov 2019 03:38:12 +0300 Subject: [PATCH] pooling update handlers, add overlay bot module --- bots/management/commands/run_bots.py | 22 +++++---- .../migrations/0006_overlaybotmoduleconfig.py | 24 ++++++++++ bots/migrations/0007_auto_20191116_0314.py | 23 +++++++++ bots/models.py | 3 +- bots/modules/__init__.py | 3 +- bots/modules/overlay.py | 48 +++++++++++++++++++ bots/templates/cabinet/bots/bot_form.html | 2 +- bots/utils.py | 6 ++- 8 files changed, 117 insertions(+), 14 deletions(-) create mode 100644 bots/migrations/0006_overlaybotmoduleconfig.py create mode 100644 bots/migrations/0007_auto_20191116_0314.py create mode 100644 bots/modules/overlay.py diff --git a/bots/management/commands/run_bots.py b/bots/management/commands/run_bots.py index 8a34315..8e67514 100644 --- a/bots/management/commands/run_bots.py +++ b/bots/management/commands/run_bots.py @@ -1,17 +1,25 @@ import logging import traceback +from multiprocessing.pool import ThreadPool import sentry_sdk from django.core.cache import cache from django.core.management import BaseCommand -from telegram import TelegramError +from telegram import TelegramError, Update from telegram.error import TimedOut +from telegram.ext import CallbackContext from bots.models import TelegramBot class Command(BaseCommand): def handle(self, *args, **options): + pool = ThreadPool(8) + + def error_handler(update: Update, ctx: CallbackContext): + sentry_sdk.capture_exception(ctx.error) + logging.exception('Exception while processing update', exc_info=ctx.error) + dispatchers = [] while True: try: @@ -20,7 +28,7 @@ class Command(BaseCommand): dispatchers = [] for bot in TelegramBot.objects.filter(active=True): try: - dispatcher = bot.build_dispatcher() + dispatcher = bot.build_dispatcher(error_handler) dispatcher.last_update_id = 0 dispatchers.append(dispatcher) except TelegramError: @@ -38,15 +46,11 @@ class Command(BaseCommand): updates = [] for update in updates: - try: - dispatcher.process_update(update) - except KeyboardInterrupt: - return - except Exception as e: - sentry_sdk.capture_exception(e) - traceback.print_exc() + pool.apply_async(dispatcher.process_update, (update,)) dispatcher.last_update_id = update.update_id + 1 except KeyboardInterrupt: + pool.terminate() + pool.join() return except Exception as e: sentry_sdk.capture_exception(e) diff --git a/bots/migrations/0006_overlaybotmoduleconfig.py b/bots/migrations/0006_overlaybotmoduleconfig.py new file mode 100644 index 0000000..9b4d0d7 --- /dev/null +++ b/bots/migrations/0006_overlaybotmoduleconfig.py @@ -0,0 +1,24 @@ +# Generated by Django 2.1.5 on 2019-11-16 00:05 + +from django.db import migrations, models + + +class Migration(migrations.Migration): + + dependencies = [ + ('bots', '0005_channelhelperbotmoduleconfig'), + ] + + operations = [ + migrations.CreateModel( + name='OverlayBotModuleConfig', + fields=[ + ('id', models.AutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')), + ('comment', models.TextField(blank=True, null=True)), + ('image', models.ImageField(upload_to='')), + ], + options={ + 'abstract': False, + }, + ), + ] diff --git a/bots/migrations/0007_auto_20191116_0314.py b/bots/migrations/0007_auto_20191116_0314.py new file mode 100644 index 0000000..8897920 --- /dev/null +++ b/bots/migrations/0007_auto_20191116_0314.py @@ -0,0 +1,23 @@ +# Generated by Django 2.1.5 on 2019-11-16 00:14 + +from django.db import migrations, models + + +class Migration(migrations.Migration): + + dependencies = [ + ('bots', '0006_overlaybotmoduleconfig'), + ] + + operations = [ + migrations.AddField( + model_name='overlaybotmoduleconfig', + name='start_text', + field=models.TextField(blank=True, null=True), + ), + migrations.AddField( + model_name='overlaybotmoduleconfig', + name='type_error_text', + field=models.TextField(blank=True, null=True), + ), + ] diff --git a/bots/models.py b/bots/models.py index fc61be5..57eb53c 100644 --- a/bots/models.py +++ b/bots/models.py @@ -23,10 +23,11 @@ class TelegramBot(models.Model): def get_bot(self): return Bot(self.bot_token) - def build_dispatcher(self): + def build_dispatcher(self, error_handler): bot = self.get_bot() bot.get_me() dispatcher = Dispatcher(bot, None, workers=0, use_context=True) + dispatcher.add_error_handler(error_handler) self.config.build_dispatcher(dispatcher) return dispatcher diff --git a/bots/modules/__init__.py b/bots/modules/__init__.py index 0126892..207bd4f 100644 --- a/bots/modules/__init__.py +++ b/bots/modules/__init__.py @@ -1,4 +1,5 @@ +from bots.modules.overlay import OverlayBotModuleConfig from .channel_helper import ChannelHelperBotModuleConfig from .echo import EchoBotModuleConfig -BOT_MODULES = [EchoBotModuleConfig, ChannelHelperBotModuleConfig] +BOT_MODULES = [EchoBotModuleConfig, ChannelHelperBotModuleConfig, OverlayBotModuleConfig] diff --git a/bots/modules/overlay.py b/bots/modules/overlay.py new file mode 100644 index 0000000..620834b --- /dev/null +++ b/bots/modules/overlay.py @@ -0,0 +1,48 @@ +import os +from tempfile import TemporaryDirectory +from time import sleep + +from PIL import Image +from django.db import models +from telegram import Update +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) + image = models.ImageField() + + MODULE_NAME = 'Overlay' + + def start_handler(self, update: Update, ctx: CallbackContext): + if 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') + out = os.path.join(d, 'out.png') + q = update.effective_message.photo[-1].get_file().download(src) + im = Image.open(q).convert('RGBA') # type: Image.Image + overlay = Image.open(self.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) + raise Exception('hi') + + 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 diff --git a/bots/templates/cabinet/bots/bot_form.html b/bots/templates/cabinet/bots/bot_form.html index dbe3607..5b88831 100644 --- a/bots/templates/cabinet/bots/bot_form.html +++ b/bots/templates/cabinet/bots/bot_form.html @@ -7,7 +7,7 @@ {% endblock %} {% block content %} -
+ {% csrf_token %}

{% if feed %}Bot "{{ feed.title }}" configuration{% else %}New bot{% endif %}

diff --git a/bots/utils.py b/bots/utils.py index 5430057..d4a4248 100644 --- a/bots/utils.py +++ b/bots/utils.py @@ -29,8 +29,10 @@ class BaseBotConfigView(CabinetViewMixin, TemplateView): def get_forms(self): bot = self.get_object() data = self.request.POST if self.request.method == 'POST' else None - return BotForm(data=data, instance=bot, module=self.get_content_type().model_class()), \ - get_config_form(self.get_content_type().model_class())(data=data, instance=bot.config if bot else None) + files = self.request.FILES if self.request.method == 'POST' else None + return BotForm(data=data, files=files, instance=bot, module=self.get_content_type().model_class()), \ + get_config_form(self.get_content_type().model_class())(data=data, files=files, + instance=bot.config if bot else None) def get_context_data(self, forms=None, **kwargs): ctx = super(BaseBotConfigView, self).get_context_data(**kwargs)