pooling update handlers, add overlay bot module

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)

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,
+            },
+        ),
+    ]

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),
+        ),
+    ]

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
 

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]

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

bots/templates/cabinet/bots/bot_form.html

@@ -7,7 +7,7 @@
 {% endblock %}
 
 {% block content %}
-    <form action="" method="post" class="card">
+    <form action="" method="post" class="card" enctype="multipart/form-data">
         {% csrf_token %}
         <header class="card-header">
             <h2 class="card-title">{% if feed %}Bot "{{ feed.title }}" configuration{% else %}New bot{% endif %}</h2>

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)