working media groups

.idea/deployment.xml

@@ -1,15 +1,6 @@
 <?xml version="1.0" encoding="UTF-8"?>
 <project version="4">
-  <component name="PublishConfigData" autoUpload="Always" serverName="bakatrouble.pw">
+  <component name="PublishConfigData" autoUpload="Always">
-    <serverData>
-      <paths name="bakatrouble.pw">
-        <serverdata>
-          <mappings>
-            <mapping deploy="/tmp/pycharm_project_802" local="$PROJECT_DIR$" />
-          </mappings>
-        </serverdata>
-      </paths>
-    </serverData>
     <option name="myAutoUpload" value="ALWAYS" />
   </component>
 </project>

main.py

@@ -6,12 +6,13 @@ from html import escape
 from queue import Queue, Empty
 from time import sleep
 from threading import Thread
-from typing import Dict
+from typing import Dict, List
 
 import sentry_sdk
 from telegram.error import Unauthorized, TelegramError
 from telegram.ext import Updater, CommandHandler, MessageHandler, Filters, CallbackQueryHandler
-from telegram import Message, Update, Bot, InlineKeyboardMarkup, InlineKeyboardButton, User
+from telegram import Message, Update, Bot, InlineKeyboardMarkup, InlineKeyboardButton, User, InputMediaPhoto, \
+    InputMediaVideo, InputMediaAnimation, InputMediaAudio, InputMediaDocument
 
 from config import BOT_TOKEN, SENTRY_DSN, MANAGEMENT_CHAT, DEBUG
 from db import get_conn, Subscriber, PersistentMapping, commit
@@ -148,6 +149,67 @@ def _sign_text(text, m: Message, limit):
     return text[:limit - len(sign)] + sign
 
 
+def _process_media_group(bot: Bot, messages: List[Message]):
+    if not messages:
+        return
+
+    m = messages[0]
+
+    current_chat = m.chat_id
+    users = conn.root.subscribers  # type: Dict[int, Subscriber]
+    if current_chat not in users:
+        if DEBUG:
+            _add_user(bot, current_chat)
+            m.reply_text('Добро пожаловать (debug)')
+        else:
+            _notify_access_request(bot, m.from_user)
+            return m.reply_text('Пожалуйста, обратитесь к @lono_contactbot')
+
+    reply_to_message_internal_id = None
+    if m.reply_to_message and m.reply_to_message.message_id in users[current_chat].messages_forward:
+        reply_to_message_internal_id = users[current_chat].messages_forward[m.reply_to_message.message_id]
+
+    media_group = []
+    for message in messages:
+        caption = _sign_text(message.caption_html, message, MAX_CAPTION_LENGTH)
+
+        if hasattr(message, 'photo') and message.photo:
+            media_group.append(InputMediaPhoto(message.photo[-1].file_id, caption=caption, parse_mode='html'))
+        elif hasattr(message, 'video') and message.video:
+            media_group.append(InputMediaVideo(message.video.file_id, caption=caption, parse_mode='html'))
+        elif hasattr(message, 'animation') and message.animation:
+            media_group.append(InputMediaAnimation(message.animation.file_id, caption=caption, parse_mode='html'))
+        elif hasattr(message, 'document') and message.document:
+            media_group.append(InputMediaDocument(message.document.file_id, caption=caption, parse_mode='html'))
+        elif hasattr(message, 'audio') and message.audio:
+            media_group.append(InputMediaAudio(message.audio.file_id, caption=caption, parse_mode='html'))
+
+    for uid, user in users.items():
+        sleep(.02)
+
+        reply_to_message_id = None
+        if reply_to_message_internal_id:
+            reply_to_message_id = user.messages_reverse.get(reply_to_message_internal_id, None)
+
+        try:
+            sent_messages = bot.send_media_group(uid, media_group, reply_to_message_id=reply_to_message_id)
+            if sent_messages:
+                user.update_from_message(sent_messages[0])
+                for r in sent_messages:
+                    user.messages_forward[r.message_id] = conn.root.counter
+                    user.messages_reverse[conn.root.counter] = r.message_id
+        except Unauthorized:
+            user = _remove_user(uid)
+            commit()
+            bot.send_message(MANAGEMENT_CHAT, f'<a href="tg://user?id={uid}">{user.name}</a> был удален '
+                                              f'из-за блокировки бота', parse_mode='html')
+        except Exception:
+            traceback.print_exc()
+            sentry_sdk.capture_exception()
+    conn.root.counter += len(messages)
+    commit()
+
+
 def _process_message(bot: Bot, m: Message):
     current_chat = m.chat_id
     users = conn.root.subscribers  # type: Dict[int, Subscriber]
@@ -241,6 +303,13 @@ def task_queue(u: Updater):
 
         try:
             m = queue.get(timeout=1)  # type: Message
+            if m.media_group_id:
+                media_group_id = m.media_group_id
+                group_messages = [m]
+                while queue.qsize() and queue.queue[0].media_group_id == media_group_id:
+                    group_messages.append(queue.get(block=False))
+                _process_media_group(u.bot, group_messages)
+            else:
                 _process_message(u.bot, m)
         except Empty:
             pass