implement reply sync (#185)

* implement reply sync #67

* bumped version from 0.2.2 to 0.2.3

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "tgcf"
-version = "0.2.2"
+version = "0.2.3"
 description = "The ultimate tool to automate custom telegram message forwarding."
 authors = ["aahnik <[email protected]>"]
 license = "MIT"

tgcf/live.py

@@ -1,41 +1,19 @@
 """The module responsible for operating tgcf in live mode."""
 
 import logging
-from typing import Dict, List
+from typing import Union
 
 from telethon import events
 from telethon.tl.custom.message import Message
 
 from tgcf import config, const
+from tgcf import storage as st
 from tgcf.bot import BOT_EVENTS
 from tgcf.plugins import apply_plugins, plugins
 from tgcf.utils import send_message
 
 
-class EventUid:
-    """The objects of this class uniquely identifies an event."""
-
-    def __init__(self, event) -> None:
-        self.chat_id = event.chat_id
-        try:
-            self.msg_id = event.id
-        except:  # pylint: disable=bare-except
-            self.msg_id = event.deleted_id
-
-    def __str__(self) -> str:
-        return f"chat={self.chat_id} msg={self.msg_id}"
-
-    def __eq__(self, other) -> bool:
-        return self.chat_id == other.chat_id and self.msg_id == other.msg_id
-
-    def __hash__(self) -> int:
-        return hash(self.__str__())
-
-
-_stored: Dict[EventUid, List[Message]] = {}
-
-
-async def new_message_handler(event) -> None:
+async def new_message_handler(event: Union[Message, events.NewMessage]) -> None:
     """Process new incoming messages."""
     chat_id = event.chat_id
 
@@ -44,30 +22,34 @@ async def new_message_handler(event) -> None:
     logging.info(f"New message received in {chat_id}")
     message = event.message
 
-    global _stored  # pylint: disable=global-statement,invalid-name
-
-    event_uid = EventUid(event)
+    event_uid = st.EventUid(event)
 
-    length = len(_stored)
+    length = len(st.stored)
     exceeding = length - const.KEEP_LAST_MANY
 
     if exceeding > 0:
-        for key in _stored:
-            del _stored[key]
+        for key in st.stored:
+            del st.stored[key]
             break
 
     to_send_to = config.from_to.get(chat_id)
 
     if to_send_to:
-        if event_uid not in _stored:
-            _stored[event_uid] = []
 
         tm = await apply_plugins(message)
         if not tm:
             return
+
+        if event.is_reply:
+            r_event = st.DummyEvent(chat_id, event.reply_to_msg_id)
+            r_event_uid = st.EventUid(r_event)
+
+        st.stored[event_uid] = {}
         for recipient in to_send_to:
+            if event.is_reply and r_event_uid in st.stored:
+                tm.reply_to = st.stored.get(r_event_uid).get(recipient)
             fwded_msg = await send_message(recipient, tm)
-            _stored[event_uid].append(fwded_msg)
+            st.stored[event_uid].update({recipient: fwded_msg})
         tm.clear()
 
 
@@ -82,17 +64,17 @@ async def edited_message_handler(event) -> None:
 
     logging.info(f"Message edited in {chat_id}")
 
-    event_uid = EventUid(event)
+    event_uid = st.EventUid(event)
 
     tm = await apply_plugins(message)
 
     if not tm:
         return
 
-    fwded_msgs = _stored.get(event_uid)
+    fwded_msgs = st.stored.get(event_uid)
 
     if fwded_msgs:
-        for msg in fwded_msgs:
+        for _, msg in fwded_msgs.items():
             if config.CONFIG.live.delete_on_edit == message.text:
                 await msg.delete()
                 await message.delete()
@@ -115,8 +97,8 @@ async def deleted_message_handler(event):
 
     logging.info(f"Message deleted in {chat_id}")
 
-    event_uid = EventUid(event)
-    fwded_msgs = _stored.get(event_uid)
+    event_uid = st.EventUid(event)
+    fwded_msgs = st.stored.get(event_uid)
     if fwded_msgs:
         for msg in fwded_msgs:
             await msg.delete()

tgcf/past.py

@@ -10,8 +10,10 @@
 
 from telethon import TelegramClient
 from telethon.errors.rpcerrorlist import FloodWaitError
+from telethon.tl.custom.message import Message
 from telethon.tl.patched import MessageService
 
+from tgcf import storage as st
 from tgcf.config import API_HASH, API_ID, CONFIG, SESSION, write_config
 from tgcf.plugins import apply_plugins
 from tgcf.utils import send_message
@@ -20,12 +22,17 @@
 async def forward_job() -> None:
     """Forward all existing messages in the concerned chats."""
     async with TelegramClient(SESSION, API_ID, API_HASH) as client:
+        client: TelegramClient
         for forward in CONFIG.forwards:
             last_id = 0
             logging.info(f"Forwarding messages from {forward.source} to {forward.dest}")
             async for message in client.iter_messages(
                 forward.source, reverse=True, offset_id=forward.offset
             ):
+                message: Message
+                event = st.DummyEvent(message.chat_id, message.id)
+                event_uid = st.EventUid(event)
+
                 if forward.end and last_id > forward.end:
                     continue
                 if isinstance(message, MessageService):
@@ -34,9 +41,18 @@ async def forward_job() -> None:
                     tm = await apply_plugins(message)
                     if not tm:
                         continue
+                    st.stored[event_uid] = {}
 
+                    if message.is_reply:
+                        r_event = st.DummyEvent(
+                            message.chat_id, message.reply_to_msg_id
+                        )
+                        r_event_uid = st.EventUid(r_event)
                     for destination in forward.dest:
-                        await send_message(destination, tm)
+                        if message.is_reply and r_event_uid in st.stored:
+                            tm.reply_to = st.stored.get(r_event_uid).get(destination)
+                        fwded_msg = await send_message(destination, tm)
+                        st.stored[event_uid].update({fwded_msg.chat_id: fwded_msg.id})
                     tm.clear()
                     last_id = message.id
                     logging.info(f"forwarding message with id = {last_id}")

tgcf/plugins/__init__.py

@@ -39,6 +39,7 @@ def __init__(self, message: Message) -> None:
         self.file_type = self.guess_file_type()
         self.new_file = None
         self.cleanup = False
+        self.reply_to = None
 
     async def get_file(self) -> str:
         """Downloads the file in the message and returns the path where its saved."""

tgcf/storage.py

@@ -0,0 +1,32 @@
+from typing import Dict
+
+from telethon.tl.custom.message import Message
+
+
+class EventUid:
+    """The objects of this class uniquely identifies a message with its chat id and message id."""
+
+    def __init__(self, event) -> None:
+        self.chat_id = event.chat_id
+        try:
+            self.msg_id = event.id
+        except:  # pylint: disable=bare-except
+            self.msg_id = event.deleted_id
+
+    def __str__(self) -> str:
+        return f"chat={self.chat_id} msg={self.msg_id}"
+
+    def __eq__(self, other) -> bool:
+        return self.chat_id == other.chat_id and self.msg_id == other.msg_id
+
+    def __hash__(self) -> int:
+        return hash(self.__str__())
+
+
+class DummyEvent:
+    def __init__(self, chat_id, msg_id):
+        self.chat_id = chat_id
+        self.id = msg_id
+
+
+stored: Dict[EventUid, Dict[int, Message]] = {}

tgcf/utils.py

@@ -22,10 +22,12 @@ async def send_message(recipient: EntityLike, tm: "TgcfMessage") -> Message:
     if CONFIG.show_forwarded_from:
         return await client.forward_messages(recipient, tm.message)
     if tm.new_file:
-        message = await client.send_file(recipient, tm.new_file, caption=tm.text)
+        message = await client.send_file(
+            recipient, tm.new_file, caption=tm.text, reply_to=tm.reply_to
+        )
         return message
     tm.message.text = tm.text
-    return await client.send_message(recipient, tm.message)
+    return await client.send_message(recipient, tm.message, reply_to=tm.reply_to)
 
 
 def cleanup(*files: str) -> None: