replies and debugging

db.py

@@ -1,3 +1,4 @@
+from BTrees.IIBTree import IIBTree
 from ZODB import DB
 from ZODB.Connection import Connection
 from ZODB.FileStorage import FileStorage
@@ -13,6 +14,15 @@ def get_conn(read_only=False) -> Connection:
     conn = db.open()
     if not hasattr(conn.root, 'subscribers'):
         conn.root.subscribers = PersistentMapping()
+    # migration 1
+    if not hasattr(conn.root, 'counter'):
+        conn.root.counter = 0
+    for user in conn.root.subscribers.values():  # type: Subscriber
+        if not hasattr(user, 'messages_forward') or not isinstance(user.messages_forward, IIBTree):
+            user.messages_forward = IIBTree()
+            user.messages_reverse = IIBTree()
+    # end migrations
+    commit()
     return conn
 
 
@@ -20,6 +30,8 @@ class Subscriber(Persistent):
     def __init__(self, user_id, name):
         self.uid = user_id
         self.name = name
+        self.messages_forward = IIBTree()
+        self.messages_reverse = IIBTree()
 
     def update_from_message(self, m: Message):
         self.name = Subscriber.get_name(m.chat)

main.py

@@ -1,17 +1,19 @@
 #!/usr/bin/env python3
 
 import logging
+import traceback
 from html import escape
 from queue import Queue, Empty
 from time import sleep
 from threading import Thread
+from typing import Dict
 
 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 config import BOT_TOKEN, SENTRY_DSN, MANAGEMENT_CHAT
+from config import BOT_TOKEN, SENTRY_DSN, MANAGEMENT_CHAT, DEBUG
 from db import get_conn, Subscriber, PersistentMapping, commit
 from send_users_list import send_users_list
 
@@ -32,6 +34,11 @@ def _notify_access_request(bot: Bot, user: User):
 
 
 def welcome(bot: Bot, update: Update):
+    if DEBUG:
+        _add_user(bot, update.effective_user.id)
+        update.message.reply_text('Добро пожаловать (debug)')
+        return
+
     if update.effective_user.id in conn.root.subscribers:
         update.message.reply_text('Вы уже являетесь участником ЛОНО')
     else:
@@ -40,12 +47,16 @@ def welcome(bot: Bot, update: Update):
 
 
 def unsubscribe(bot: Bot, update: Update):
-    del conn.root.subscribers[update.message.chat_id]
+    user = _remove_user(update.message.chat_id)
-    commit()
     update.message.reply_text('Вы были отписаны от бота. '
                               'Обратитесь к @lono_contactbot если вы хотите подписаться снова.')
-    user = update.message.from_user
+    bot.send_message(MANAGEMENT_CHAT, f'<a href="tg://user?id={user.id}">{escape(user.name)}</a> отписался')
-    bot.send_message(MANAGEMENT_CHAT, f'<a href="tg://user?id={user.id}">{escape(user.full_name)}</a> отписался')
+
+
+def _add_user(bot, uid):
+    user = conn.root.subscribers[uid] = Subscriber.from_chat(bot.get_chat(uid))
+    commit()
+    return user
 
 
 def add_user(bot: Bot, update: Update, groups=(), args=()):
@@ -69,8 +80,7 @@ def add_user(bot: Bot, update: Update, groups=(), args=()):
         pass
 
     try:
-        user = conn.root.subscribers[uid] = Subscriber.from_chat(bot.get_chat(uid))
+        user = _add_user(bot, uid)
-        commit()
         if update.callback_query:
             update.callback_query.message.edit_reply_markup()
         bot.send_message(MANAGEMENT_CHAT, f'<a href="tg://user?id={uid}">{escape(user.name)}</a> был добавлен',
@@ -80,6 +90,13 @@ def add_user(bot: Bot, update: Update, groups=(), args=()):
         bot.send_message(MANAGEMENT_CHAT, str(e))
 
 
+def _remove_user(uid):
+    user = conn.root.subscribers[uid]
+    del conn.root.subscribers[uid]
+    commit()
+    return user
+
+
 def remove_user(bot: Bot, update: Update, groups=(), args=()):
     if update.callback_query:
         update.callback_query.answer()
@@ -101,10 +118,8 @@ def remove_user(bot: Bot, update: Update, groups=(), args=()):
         pass
 
     try:
-        name = conn.root.subscribers[uid].name
+        user = _remove_user(uid)
-        del conn.root.subscribers[uid]
+        bot.send_message(MANAGEMENT_CHAT, f'<a href="tg://user?id={uid}">{escape(user.name)}</a> был удален',
-        commit()
-        bot.send_message(MANAGEMENT_CHAT, f'<a href="tg://user?id={uid}">{escape(name)}</a> был удален',
                          parse_mode='html')
         if update.callback_query:
             update.callback_query.message.edit_reply_markup()
@@ -134,66 +149,87 @@ def _sign_text(text, m: Message, limit):
 
 def _process_message(bot: Bot, m: Message):
     current_chat = m.chat_id
-    users = conn.root.subscribers  # type: PersistentMapping
+    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')
 
     text = _sign_text(m.text_html, m, MAX_MESSAGE_LENGTH)
     caption = _sign_text(m.caption_html, m, MAX_CAPTION_LENGTH)
 
+    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]
+
     for uid, user in users.items():
-        if uid == current_chat:
-            continue
         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:
             r = None
             if m.forward_date:
                 r = m.forward(uid)
             elif hasattr(m, 'audio') and m.audio:
                 a = m.audio
-                r = bot.send_audio(uid, a.file_id, a.duration, a.performer, a.title, caption, parse_mode='html')
+                r = bot.send_audio(uid, a.file_id, a.duration, a.performer, a.title, caption,
+                                   reply_to_message_id=reply_to_message_id, parse_mode='html')
             elif hasattr(m, 'document') and m.document:
                 d = m.document
-                r = bot.send_document(uid, d.file_id, d.file_name, caption, parse_mode='html')
+                r = bot.send_document(uid, d.file_id, d.file_name, caption, reply_to_message_id=reply_to_message_id,
+                                      parse_mode='html')
             elif hasattr(m, 'photo') and m.photo:
                 p = m.photo
-                r = bot.send_photo(uid, p[-1].file_id, caption, parse_mode='html')
+                r = bot.send_photo(uid, p[-1].file_id, caption, reply_to_message_id=reply_to_message_id,
+                                   parse_mode='html')
             elif hasattr(m, 'sticker') and m.sticker:
                 s = m.sticker
-                r = bot.send_sticker(uid, s.file_id)
+                r = bot.send_sticker(uid, s.file_id, reply_to_message_id=reply_to_message_id)
             elif hasattr(m, 'video') and m.video:
                 v = m.video
-                r = bot.send_video(uid, v.file_id, v.duration, caption, parse_mode='html')
+                r = bot.send_video(uid, v.file_id, v.duration, caption, reply_to_message_id=reply_to_message_id,
+                                   parse_mode='html')
             elif hasattr(m, 'voice') and m.voice:
                 v = m.voice
-                r = bot.send_voice(uid, v.file_id, v.duration, caption, parse_mode='html')
+                r = bot.send_voice(uid, v.file_id, v.duration, caption, reply_to_message_id=reply_to_message_id,
+                                   parse_mode='html')
             elif hasattr(m, 'video_note') and m.video_note:
                 vn = m.video_note
-                r = bot.send_video_note(uid, vn.file_id, vn.duration, vn.length)
+                r = bot.send_video_note(uid, vn.file_id, vn.duration, vn.length,
+                                        reply_to_message_id=reply_to_message_id)
             elif hasattr(m, 'contact') and m.contact:
                 c = m.contact
-                r = bot.send_contact(uid, c.phone_number, c.first_name, c.last_name)
+                r = bot.send_contact(uid, c.phone_number, c.first_name, c.last_name,
+                                     reply_to_message_id=reply_to_message_id)
             elif hasattr(m, 'location') and m.location:
                 l = m.location
-                r = bot.send_location(uid, l.latitude, l.longitude)
+                r = bot.send_location(uid, l.latitude, l.longitude, reply_to_message_id=reply_to_message_id)
             elif hasattr(m, 'venue') and m.venue:
                 v = m.venue
                 l = v.location
-                r = bot.send_venue(uid, l.latitude, l.longitude, v.title, v.address, v.foursquare_id)
+                r = bot.send_venue(uid, l.latitude, l.longitude, v.title, v.address, v.foursquare_id,
+                                   reply_to_message_id=reply_to_message_id)
             elif hasattr(m, 'text') and m.text:
-                r = bot.send_message(uid, text, 'html')
+                r = bot.send_message(uid, text, 'html', reply_to_message_id=reply_to_message_id)
             if r:
                 user.update_from_message(r)
+                user.messages_forward[r.message_id] = conn.root.counter
+                user.messages_reverse[conn.root.counter] = r.message_id
         except Unauthorized:
-            name = conn.root.subscribers[uid].name
+            user = _remove_user(uid)
-            del conn.root.subscribers[uid]
             commit()
-            bot.send_message(MANAGEMENT_CHAT, f'<a href="tg://user?id={uid}">{name}</a> был удален '
+            bot.send_message(MANAGEMENT_CHAT, f'<a href="tg://user?id={uid}">{user.name}</a> был удален '
                                               f'из-за блокировки бота', parse_mode='html')
-        except TelegramError:
+        except Exception:
+            traceback.print_exc()
             sentry_sdk.capture_exception()
-
+    conn.root.counter += 1
     commit()
 
 
@@ -208,6 +244,7 @@ def task_queue(u: Updater):
         except Empty:
             pass
         except:
+            traceback.print_exc()
             sentry_sdk.capture_exception()