From f7754841723211a3ed013b772cf49172ba7b2b9c Mon Sep 17 00:00:00 2001 From: Lonami Exo Date: Mon, 24 Jan 2022 11:05:58 +0100 Subject: [PATCH] Properly load and save channel state --- telethon/_client/telegrambaseclient.py | 11 ++++++++--- telethon/_updates/entitycache.py | 3 +++ 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/telethon/_client/telegrambaseclient.py b/telethon/_client/telegrambaseclient.py index 5c32d086..480263e4 100644 --- a/telethon/_client/telegrambaseclient.py +++ b/telethon/_client/telegrambaseclient.py @@ -14,7 +14,7 @@ from .._crypto import rsa from .._misc import markdown, enums, helpers from .._network import MTProtoSender, Connection, transports from .._sessions import Session, SQLiteSession, MemorySession -from .._sessions.types import DataCenter, SessionState +from .._sessions.types import DataCenter, SessionState, EntityType, ChannelState from .._updates import EntityCache, MessageBox DEFAULT_DC_ID = 2 @@ -235,7 +235,12 @@ async def connect(self: 'TelegramClient') -> None: else: try_fetch_user = self._session_state.user_id == 0 if self._catch_up: - self._message_box.load(self._session_state, await self._session.get_all_channel_states()) + channel_states = await self._session.get_all_channel_states() + self._message_box.load(self._session_state, channel_states) + for state in channel_states: + entity = await self._session.get_entity(EntityType.CHANNEL, state.channel_id) + if entity: + self._entity_cache.put(entity) dc = all_dcs.get(self._session_state.dc_id) if dc is None: @@ -366,7 +371,7 @@ async def _disconnect(self: 'TelegramClient'): session_state, channel_states = self._message_box.session_state() for channel_id, pts in channel_states.items(): - await self._session.insert_channel_state(channel_id, pts) + await self._session.insert_channel_state(ChannelState(channel_id=channel_id, pts=pts)) await self._replace_session_state(**session_state) diff --git a/telethon/_updates/entitycache.py b/telethon/_updates/entitycache.py index ce89eb4f..8dc95693 100644 --- a/telethon/_updates/entitycache.py +++ b/telethon/_updates/entitycache.py @@ -98,3 +98,6 @@ class EntityCache: def get_all_entities(self): return [Entity(ty, id, hash) for id, (hash, ty) in self.hash_map.items()] + + def put(self, entity): + self.hash_map[entity.id] = (entity.access_hash, entity.ty)