diff --git a/stdplugins/chain.py b/stdplugins/chain.py index 8c28ffc..f35bfff 100644 --- a/stdplugins/chain.py +++ b/stdplugins/chain.py @@ -2,22 +2,48 @@ # License, v. 2.0. If a copy of the MPL was not distributed with this # file, You can obtain one at http://mozilla.org/MPL/2.0/. +from collections import defaultdict + from telethon import events from telethon.tl.functions.messages import SaveDraftRequest +def intify(d): + for k, v in d.items(): + if isinstance(v, dict): + v = dict(intify(v)) + yield int(k), v + +global_cache = defaultdict(lambda: {}, intify(storage.cache or {})) + @borg.on(events.NewMessage(pattern=r"\.chain", outgoing=True)) async def _(event): - await event.edit("Counting...") - count = -1 - message = event.message - while message: - reply = await message.get_reply_message() + cache = global_cache[event.chat_id] + + message = event.reply_to_msg_id + count = 0 + while message is not None: + reply = cache.get(message, -1) + if reply == -1: + reply = None + async for m in borg.iter_messages( + event.chat_id, + ids=message + ): + reply = m.reply_to_msg_id + cache[message] = reply + if len(cache) % 10 == 0: + await event.edit(f"Counting... ({len(cache)} cached entries)") + if reply is None: await borg(SaveDraftRequest( await event.get_input_chat(), "", - reply_to_msg_id=message.id + reply_to_msg_id=message )) + message = reply count += 1 + + if count: + storage.cache = global_cache await event.edit(f"Chain length: {count}")