Fix pts from channels is different (#1160)

2025-08-08 21:10:29 +00:00 · 2019-04-21 13:56:14 +02:00
parent 8edbfbdced
commit c1880c9191
4 changed files with 141 additions and 50 deletions
--- a/telethon/statecache.py
+++ b/telethon/statecache.py
@@ -0,0 +1,117 @@
+import datetime
+
+from .tl import types
+
+
+class StateCache:
+    """
+    In-memory update state cache, defaultdict-like behaviour.
+    """
+    def __init__(self, initial):
+        # We only care about the pts and the date. By using a tuple which
+        # is lightweight and immutable we can easily copy them around to
+        # each update in case they need to fetch missing entities.
+        if initial:
+            self._pts_date = initial.pts, initial.date
+        else:
+            self._pts_date = 1, datetime.datetime.now()
+
+    def reset(self):
+        self.__dict__.clear()
+        self._pts_date = (1, 1)
+
+    # TODO Call this when receiving responses too...?
+    def update(
+            self,
+            update,
+            *,
+            channel_id=None,
+            has_pts=(
+                types.UpdateNewMessage,
+                types.UpdateDeleteMessages,
+                types.UpdateReadHistoryInbox,
+                types.UpdateReadHistoryOutbox,
+                types.UpdateWebPage,
+                types.UpdateReadMessagesContents,
+                types.UpdateEditMessage,
+                types.updates.State,
+                types.updates.DifferenceTooLong,
+                types.UpdateShortMessage,
+                types.UpdateShortChatMessage,
+                types.UpdateShortSentMessage
+            ),
+            has_date=(
+                types.UpdateUserPhoto,
+                types.UpdateEncryption,
+                types.UpdateEncryptedMessagesRead,
+                types.UpdateChatParticipantAdd,
+                types.updates.DifferenceEmpty,
+                types.UpdateShortMessage,
+                types.UpdateShortChatMessage,
+                types.UpdateShort,
+                types.UpdatesCombined,
+                types.Updates,
+                types.UpdateShortSentMessage,
+            ),
+            has_channel_pts=(
+                types.UpdateChannelTooLong,
+                types.UpdateNewChannelMessage,
+                types.UpdateDeleteChannelMessages,
+                types.UpdateEditChannelMessage,
+                types.UpdateChannelWebPage,
+                types.updates.ChannelDifferenceEmpty,
+                types.updates.ChannelDifferenceTooLong,
+                types.updates.ChannelDifference
+            )
+    ):
+        """
+        Update the state with the given update.
+        """
+        has_pts = isinstance(update, has_pts)
+        has_date = isinstance(update, has_date)
+        has_channel_pts = isinstance(update, has_channel_pts)
+        if has_pts and has_date:
+            self._pts_date = update.pts, update.date
+        elif has_pts:
+            self._pts_date = update.pts, self._pts_date[1]
+        elif has_date:
+            self._pts_date = self._pts_date[0], update.date
+
+        if has_channel_pts:
+            if channel_id is None:
+                channel_id = self.get_channel_id(update)
+
+            if channel_id is None:
+                pass  # TODO log, but shouldn't happen
+            else:
+                self.__dict__[channel_id] = update.pts
+
+    @staticmethod
+    def get_channel_id(
+            update,
+            has_channel_id=(
+                types.UpdateChannelTooLong,
+                types.UpdateDeleteChannelMessages,
+                types.UpdateChannelWebPage
+            ),
+            has_message=(
+                types.UpdateNewChannelMessage,
+                types.UpdateEditChannelMessage
+            )
+    ):
+        # Will only fail for *difference, where channel_id is known
+        if isinstance(update, has_channel_id):
+            return update.channel_id
+        elif isinstance(update, has_message):
+            return update.message.to_id.channel_id
+        else:
+            return None
+
+    def __getitem__(self, item):
+        """
+        Gets the corresponding ``(pts, date)`` for the given ID or peer,
+        """
+        if item is None:
+            return self._pts_date
+        else:
+            return self.__dict__.get(item, 1)