Make dialogs and contacts fetching smarter

This commit is contained in:
Dan 2018-04-13 15:47:07 +02:00
parent d9aa2a3ab0
commit 973cb4b2b6
2 changed files with 47 additions and 42 deletions

View File

@ -122,6 +122,7 @@ class Client:
DIALOGS_AT_ONCE = 100 DIALOGS_AT_ONCE = 100
UPDATES_WORKERS = 1 UPDATES_WORKERS = 1
DOWNLOAD_WORKERS = 1 DOWNLOAD_WORKERS = 1
OFFLINE_SLEEP = 900
def __init__(self, def __init__(self,
session_name: str, session_name: str,
@ -218,8 +219,14 @@ class Client:
self.save_session() self.save_session()
if self.token is None: if self.token is None:
now = time.time()
if abs(now - self.date) > Client.OFFLINE_SLEEP:
self.get_dialogs() self.get_dialogs()
self.get_contacts() self.get_contacts()
else:
self.send(functions.messages.GetPinnedDialogs())
self.get_dialogs_chunk(0)
else: else:
self.send(functions.updates.GetState()) self.send(functions.updates.GetState())
@ -845,23 +852,29 @@ class Client:
s = json.load(f) s = json.load(f)
except FileNotFoundError: except FileNotFoundError:
self.dc_id = 1 self.dc_id = 1
self.date = int(time.time()) self.date = 0
self.auth_key = Auth(self.dc_id, self.test_mode, self.proxy).create() self.auth_key = Auth(self.dc_id, self.test_mode, self.proxy).create()
else: else:
self.dc_id = s["dc_id"] self.dc_id = s["dc_id"]
self.test_mode = s["test_mode"] self.test_mode = s["test_mode"]
self.auth_key = base64.b64decode("".join(s["auth_key"])) self.auth_key = base64.b64decode("".join(s["auth_key"]))
self.user_id = s["user_id"] self.user_id = s["user_id"]
self.date = s.get("date", int(time.time())) self.date = s.get("date", 0)
for k, v in s.get("peers_by_id", {}).items(): for k, v in s.get("peers_by_id", {}).items():
self.peers_by_id[int(k)] = utils.get_input_peer(int(k), v) self.peers_by_id[int(k)] = utils.get_input_peer(int(k), v)
for k, v in s.get("peers_by_username", {}).items(): for k, v in s.get("peers_by_username", {}).items():
self.peers_by_username[k] = self.peers_by_id[v] peer = self.peers_by_id.get(v, None)
if peer:
self.peers_by_username[k] = peer
for k, v in s.get("peers_by_phone", {}).items(): for k, v in s.get("peers_by_phone", {}).items():
self.peers_by_phone[k] = self.peers_by_id[v] peer = self.peers_by_id.get(v, None)
if peer:
self.peers_by_phone[k] = peer
def save_session(self): def save_session(self):
auth_key = base64.b64encode(self.auth_key).decode() auth_key = base64.b64encode(self.auth_key).decode()
@ -880,53 +893,34 @@ class Client:
indent=4 indent=4
) )
def get_dialogs(self): def get_dialogs_chunk(self, offset_date):
def parse_dialogs(d): r = self.send(
for m in reversed(d.messages):
if isinstance(m, types.MessageEmpty):
continue
else:
return m.date
else:
return 0
pinned_dialogs = self.send(functions.messages.GetPinnedDialogs())
parse_dialogs(pinned_dialogs)
dialogs = self.send(
functions.messages.GetDialogs(
0, 0, types.InputPeerEmpty(),
self.DIALOGS_AT_ONCE, True
)
)
offset_date = parse_dialogs(dialogs)
log.info("Entities count: {}".format(len(self.peers_by_id)))
while len(dialogs.dialogs) == self.DIALOGS_AT_ONCE:
try:
dialogs = self.send(
functions.messages.GetDialogs( functions.messages.GetDialogs(
offset_date, 0, types.InputPeerEmpty(), offset_date, 0, types.InputPeerEmpty(),
self.DIALOGS_AT_ONCE, True self.DIALOGS_AT_ONCE, True
) )
) )
log.info("Total peers: {}".format(len(self.peers_by_id)))
return r
def get_dialogs(self):
self.send(functions.messages.GetPinnedDialogs())
dialogs = self.get_dialogs_chunk(0)
offset_date = utils.get_offset_date(dialogs)
while len(dialogs.dialogs) == self.DIALOGS_AT_ONCE:
try:
dialogs = self.get_dialogs_chunk(offset_date)
except FloodWait as e: except FloodWait as e:
log.warning("get_dialogs flood: waiting {} seconds".format(e.x)) log.warning("get_dialogs flood: waiting {} seconds".format(e.x))
time.sleep(e.x) time.sleep(e.x)
continue continue
offset_date = parse_dialogs(dialogs) offset_date = utils.get_offset_date(dialogs)
log.info("Entities count: {}".format(len(self.peers_by_id)))
self.send( self.get_dialogs_chunk(0)
functions.messages.GetDialogs(
0, 0, types.InputPeerEmpty(),
self.DIALOGS_AT_ONCE, True
)
)
log.info("Entities count: {}".format(len(self.peers_by_id)))
def resolve_peer(self, peer_id: int or str): def resolve_peer(self, peer_id: int or str):
"""Use this method to get the *InputPeer* of a known *peer_id*. """Use this method to get the *InputPeer* of a known *peer_id*.
@ -2850,7 +2844,7 @@ class Client:
continue continue
else: else:
if isinstance(contacts, types.contacts.Contacts): if isinstance(contacts, types.contacts.Contacts):
log.info("Contacts count: {}".format(len(contacts.users))) log.info("Total contacts: {}".format(len(self.peers_by_phone)))
return contacts return contacts

View File

@ -9,9 +9,20 @@ def get_peer_id(input_peer) -> int:
) )
def get_input_peer(peer_id: int, access_hash: int = 0): def get_input_peer(peer_id: int, access_hash: int):
return ( return (
types.InputPeerUser(peer_id, access_hash) if peer_id > 0 types.InputPeerUser(peer_id, access_hash) if peer_id > 0
else types.InputPeerChannel(int(str(peer_id)[4:]), access_hash) if str(peer_id).startswith("-100") else types.InputPeerChannel(int(str(peer_id)[4:]), access_hash)
if (str(peer_id).startswith("-100") and access_hash)
else types.InputPeerChat(-peer_id) else types.InputPeerChat(-peer_id)
) )
def get_offset_date(dialogs):
for m in reversed(dialogs.messages):
if isinstance(m, types.MessageEmpty):
continue
else:
return m.date
else:
return 0