diff --git a/bots/migrations/0009_channelhelperbotmoduleconfig_photo_group_size_and_more.py b/bots/migrations/0009_channelhelperbotmoduleconfig_photo_group_size_and_more.py new file mode 100644 index 0000000..6100a21 --- /dev/null +++ b/bots/migrations/0009_channelhelperbotmoduleconfig_photo_group_size_and_more.py @@ -0,0 +1,28 @@ +# Generated by Django 4.1 on 2022-08-25 06:44 + +from django.db import migrations, models + + +class Migration(migrations.Migration): + + dependencies = [ + ("bots", "0008_queueditem_processed"), + ] + + operations = [ + migrations.AddField( + model_name="channelhelperbotmoduleconfig", + name="photo_group_size", + field=models.PositiveIntegerField(default=10), + ), + migrations.AddField( + model_name="channelhelperbotmoduleconfig", + name="send_photo_groups", + field=models.BooleanField(default=False), + ), + migrations.AddField( + model_name="channelhelperbotmoduleconfig", + name="send_photo_groups_threshold", + field=models.PositiveIntegerField(default=0, help_text="0 = disabled"), + ), + ] diff --git a/bots/modules/channel_helper.py b/bots/modules/channel_helper.py index 8f2d9ff..b2493fa 100644 --- a/bots/modules/channel_helper.py +++ b/bots/modules/channel_helper.py @@ -9,7 +9,7 @@ import imagehash import requests from PIL import Image from django.db import models -from telegram import Update, Bot +from telegram import Update, Bot, InputMediaPhoto from telegram.ext import Dispatcher, CallbackContext, MessageHandler, Filters, CommandHandler from jsonrpc import Dispatcher as RPCDispatcher from djconfig import config @@ -21,6 +21,9 @@ class ChannelHelperBotModuleConfig(TelegramBotModuleConfig): chat_id = models.CharField(max_length=32) queued = models.BooleanField(default=False) users = models.ManyToManyField(BotUser) + send_photo_groups = models.BooleanField(default=False) + send_photo_groups_threshold = models.PositiveIntegerField(default=0, help_text='0 = disabled') + photo_group_size = models.PositiveIntegerField(default=10) MODULE_NAME = 'Channel helper' @@ -158,6 +161,19 @@ class ChannelHelperBotModuleConfig(TelegramBotModuleConfig): dispatcher.add_handler(MessageHandler(Filters.private, self.handle_message)) return dispatcher + def should_send_photo_group(self): + # noinspection PyChainedComparisons + return self.send_photo_groups or ( + self.send_photo_groups_threshold > 0 and + self.queued_items.filter(type='photo', processed=False).count() > self.send_photo_groups_threshold + ) + + def get_additional_images(self, for_id: int): + return self.queued_items\ + .filter(type='photo', processed=False)\ + .exclude(pk=for_id)\ + .order_by('?')[:self.photo_group_size - 1] + class QueuedItem(models.Model): config = models.ForeignKey(ChannelHelperBotModuleConfig, on_delete=models.CASCADE, related_name='queued_items') @@ -168,6 +184,13 @@ class QueuedItem(models.Model): processed = models.BooleanField(default=False) def send(self, bot: Bot): - getattr(bot, 'send_' + self.type)(self.config.chat_id, *json.loads(self.args)) - self.processed = True - self.save() + if self.type == 'photo' and self.config.should_send_photo_group(): + items = [self] + list(self.config.get_additional_images(self.pk)) + bot.send_media_group(self.config.chat_id, [InputMediaPhoto(json.loads(i.args)[0]) for i in items]) + for i in items: + i.processed = True + i.save() + else: + getattr(bot, 'send_' + self.type)(self.config.chat_id, *json.loads(self.args)) + self.processed = True + self.save()