LonamiWebs/Telethon
1
1
Fork 0
mirror of https://github.com/LonamiWebs/Telethon.git synced 2024-09-21 03:08:49 +03:00
Code Issues Packages Projects Releases Wiki Activity

Support async def in sessions (#1013)

Browse Source
This commit is contained in:
josephbiko 2018-10-05 20:25:49 +02:00 committed by Lonami
parent 653f3c043d
commit 3dd8b7c6d1
10 changed files with 49 additions and 28 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

0
.gitignore vendored Executable file → Normal file
View File

2
telethon/client/auth.py
View File

@ -399,7 +399,7 @@ class AuthMethods(MessageParseMethods, UserMethods):
return False
await self.disconnect()
self.session.delete()
await self.session.delete()
self._authorized = False
return True

4
telethon/client/downloads.py
View File

@ -223,9 +223,9 @@ class DownloadMethods(UserMethods):
config = await self(functions.help.GetConfigRequest())
for option in config.dc_options:
if option.ip_address == self.session.server_address:
self.session.set_dc(
await self.session.set_dc(
option.id, option.ip_address, option.port)
self.session.save()
await self.session.save()
break
# TODO Figure out why the session may have the wrong DC ID

2
telethon/client/messages.py
View File

@ -662,7 +662,7 @@ class MessageMethods(UploadMethods, ButtonMethods, MessageParseMethods):
reply_markup=self.build_reply_markup(buttons)
)
msg = self._get_response_message(request, await self(request), entity)
self._cache_media(msg, file, file_handle)
await self._cache_media(msg, file, file_handle)
return msg
async def delete_messages(self, entity, message_ids, *, revoke=True):

19
telethon/client/telegrambaseclient.py
View File

@ -14,6 +14,7 @@ from ..network import MTProtoSender, ConnectionTcpFull
from ..sessions import Session, SQLiteSession, MemorySession
from ..tl import TLObject, functions, types
from ..tl.alltlobjects import LAYER
from ..utils import AsyncClassWrapper
DEFAULT_DC_ID = 4
DEFAULT_IPV4_IP = '149.154.167.51'
@ -196,7 +197,7 @@ class TelegramBaseClient(abc.ABC):
)
self.flood_sleep_threshold = flood_sleep_threshold
self.session = session
self.session = AsyncClassWrapper(session)
self.api_id = int(api_id)
self.api_hash = api_hash
@ -327,8 +328,8 @@ class TelegramBaseClient(abc.ABC):
await self._disconnect()
if getattr(self, 'session', None):
if getattr(self, '_state', None):
self.session.set_update_state(0, self._state)
self.session.close()
await self.session.set_update_state(0, self._state)
await self.session.close()
async def _disconnect(self):
"""
@ -366,22 +367,22 @@ class TelegramBaseClient(abc.ABC):
__log__.info('Reconnecting to new data center %s', new_dc)
dc = await self._get_dc(new_dc)
self.session.set_dc(dc.id, dc.ip_address, dc.port)
await self.session.set_dc(dc.id, dc.ip_address, dc.port)
# auth_key's are associated with a server, which has now changed
# so it's not valid anymore. Set to None to force recreating it.
self.session.auth_key = None
self.session.save()
await self.session.save()
await self._disconnect()
return await self.connect()
def _auth_key_callback(self, auth_key):
async def _auth_key_callback(self, auth_key):
"""
Callback from the sender whenever it needed to generate a
new authorization key. This means we are not authorized.
"""
self._authorized = None
self.session.auth_key = auth_key
self.session.save()
await self.session.save()
# endregion
@ -469,8 +470,8 @@ class TelegramBaseClient(abc.ABC):
session = self._exported_sessions.get(cdn_redirect.dc_id)
if not session:
dc = await self._get_dc(cdn_redirect.dc_id, cdn=True)
session = self.session.clone()
session.set_dc(dc.id, dc.ip_address, dc.port)
session = await self.session.clone()
await session.set_dc(dc.id, dc.ip_address, dc.port)
self._exported_sessions[cdn_redirect.dc_id] = session
__log__.info('Creating new CDN client')

10
telethon/client/updates.py
View File

@ -137,7 +137,7 @@ class UpdateMethods(UserMethods):
This can also be used to forcibly fetch new updates if there are any.
"""
state = self.session.get_update_state(0)
state = await self.session.get_update_state(0)
if not state or not state.pts:
state = await self(functions.updates.GetStateRequest())
@ -172,15 +172,15 @@ class UpdateMethods(UserMethods):
state.pts = d.pts
break
finally:
self.session.set_update_state(0, state)
await self.session.set_update_state(0, state)
self.session.catching_up = False
# endregion
# region Private methods
def _handle_update(self, update):
self.session.process_entities(update)
async def _handle_update(self, update):
await self.session.process_entities(update)
if isinstance(update, (types.Updates, types.UpdatesCombined)):
entities = {utils.get_peer_id(x): x for x in
itertools.chain(update.users, update.chats)}
@ -236,7 +236,7 @@ class UpdateMethods(UserMethods):
# inserted because this is a rather expensive operation
# (default's sqlite3 takes ~0.1s to commit changes). Do
# it every minute instead. No-op if there's nothing new.
self.session.save()
await self.session.save()
# We need to send some content-related request at least hourly
# for Telegram to keep delivering updates, otherwise they will

10
telethon/client/uploads.py
View File

@ -174,7 +174,7 @@ class UploadMethods(ButtonMethods, MessageParseMethods, UserMethods):
entities=msg_entities, reply_markup=markup, silent=silent
)
msg = self._get_response_message(request, await self(request), entity)
self._cache_media(msg, file, file_handle, force_document=force_document)
await self._cache_media(msg, file, file_handle, force_document=force_document)
return msg
@ -211,7 +211,7 @@ class UploadMethods(ButtonMethods, MessageParseMethods, UserMethods):
entity, media=types.InputMediaUploadedPhoto(fh)
))
input_photo = utils.get_input_photo(r.photo)
self.session.cache_file(fh.md5, fh.size, input_photo)
await self.session.cache_file(fh.md5, fh.size, input_photo)
fh = input_photo
if captions:
@ -326,7 +326,7 @@ class UploadMethods(ButtonMethods, MessageParseMethods, UserMethods):
file = stream.read()
hash_md5.update(file)
if use_cache:
cached = self.session.get_file(
cached = await self.session.get_file(
hash_md5.digest(), file_size, cls=use_cache
)
if cached:
@ -446,7 +446,7 @@ class UploadMethods(ButtonMethods, MessageParseMethods, UserMethods):
)
return file_handle, media
def _cache_media(self, msg, file, file_handle,
async def _cache_media(self, msg, file, file_handle,
force_document=False):
if file and msg and isinstance(file_handle,
custom.InputSizedFile):
@ -457,6 +457,6 @@ class UploadMethods(ButtonMethods, MessageParseMethods, UserMethods):
to_cache = utils.get_input_photo(msg.media.photo)
else:
to_cache = utils.get_input_document(msg.media.document)
self.session.cache_file(md5, size, to_cache)
await self.session.cache_file(md5, size, to_cache)
# endregion

8
telethon/client/users.py
View File

@ -48,7 +48,7 @@ class UserMethods(TelegramBaseClient):
exceptions.append(e)
results.append(None)
continue
self.session.process_entities(result)
await self.session.process_entities(result)
exceptions.append(None)
results.append(result)
request_index += 1
@ -58,7 +58,7 @@ class UserMethods(TelegramBaseClient):
return results
else:
result = await future
self.session.process_entities(result)
await self.session.process_entities(result)
return result
except (errors.ServerError, errors.RpcCallFailError) as e:
__log__.warning('Telegram is having internal issues %s: %s',
@ -288,7 +288,7 @@ class UserMethods(TelegramBaseClient):
try:
# First try to get the entity from cache, otherwise figure it out
return self.session.get_input_entity(peer)
return await self.session.get_input_entity(peer)
except ValueError:
pass
@ -393,7 +393,7 @@ class UserMethods(TelegramBaseClient):
try:
# Nobody with this username, maybe it's an exact name/title
return await self.get_entity(
self.session.get_input_entity(string))
await self.session.get_input_entity(string))
except ValueError:
pass

3
telethon/network/mtprotosender.py
View File

@ -21,6 +21,7 @@ from ..tl.types import (
MsgNewDetailedInfo, NewSessionCreated, MsgDetailedInfo, MsgsStateReq,
MsgsStateInfo, MsgsAllInfo, MsgResendReq, upload
)
from ..utils import AsyncClassWrapper
__log__ = logging.getLogger(__name__)
@ -213,7 +214,7 @@ class MTProtoSender:
await authenticator.do_authentication(plain)
if self._auth_key_callback:
self._auth_key_callback(state.auth_key)
await self._auth_key_callback(state.auth_key)
break
except (SecurityError, AssertionError) as e:

19
telethon/utils.py
View File

@ -16,6 +16,7 @@ from types import GeneratorType
from .extensions import markdown, html
from .helpers import add_surrogate, del_surrogate
from .tl import types
import inspect
try:
import hachoir
@ -994,3 +995,21 @@ def get_appropriated_part_size(file_size):
return 512
raise ValueError('File size too large')
class AsyncClassWrapper:
def __init__(self, wrapped):
self.wrapped = wrapped
def __getattr__(self, item):
w = getattr(self.wrapped, item)
async def wrapper(*args, **kwargs):
val = w(*args, **kwargs)
return await val if inspect.isawaitable(val) else val
if callable(w):
return wrapper
elif isinstance(w, property):
return w.fget(self.wrapped)
else:
return w