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

Add support to finish takeouts automatically by the context manager

Browse Source
This commit is contained in:
Dmitry D. Chernov 2019-02-10 02:30:14 +10:00
parent 78b6b02342
commit c848c23b75
1 changed files with 29 additions and 17 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

46
telethon/client/account.py
View File

@ -12,13 +12,14 @@ class _TakeoutClient:
""" """
__PROXY_INTERFACE = ('__enter__', '__exit__', '__aenter__', '__aexit__') __PROXY_INTERFACE = ('__enter__', '__exit__', '__aenter__', '__aexit__')
def __init__(self, client, request): def __init__(self, finalize, client, request):
# We use the name mangling for attributes to make them inaccessible # We use the name mangling for attributes to make them inaccessible
# from within the shadowed client object and to distinguish them from # from within the shadowed client object and to distinguish them from
# its own attributes where needed. # its own attributes where needed.
self.__finalize = finalize
self.__client = client self.__client = client
self.__request = request self.__request = request
self.__success = True self.__success = None
@property @property
def success(self): def success(self):
@ -50,9 +51,16 @@ class _TakeoutClient:
def __exit__(self, *args): def __exit__(self, *args):
return self.__client.loop.run_until_complete(self.__aexit__(*args)) return self.__client.loop.run_until_complete(self.__aexit__(*args))
async def __aexit__(self, *args): async def __aexit__(self, exc_type, exc_value, traceback):
if self.__success is None and self.__finalize:
self.__success = exc_type is None
if self.__success is not None: if self.__success is not None:
await self.__client.end_takeout(self.__success) result = await self(functions.account.FinishTakeoutSessionRequest(
self.__success))
if not result:
raise ValueError("Failed to finish the takeout.")
self.session.takeout_id = None
async def __call__(self, request, ordered=False): async def __call__(self, request, ordered=False):
takeout_id = self.__client.session.takeout_id takeout_id = self.__client.session.takeout_id
@ -104,8 +112,8 @@ class _TakeoutClient:
class AccountMethods(UserMethods): class AccountMethods(UserMethods):
def takeout( def takeout(
self, *, contacts=None, users=None, chats=None, megagroups=None, self, finalize=True, *, contacts=None, users=None, chats=None,
channels=None, files=None, max_file_size=None): megagroups=None, channels=None, files=None, max_file_size=None):
""" """
Creates a proxy object over the current :ref:`TelegramClient` through Creates a proxy object over the current :ref:`TelegramClient` through
which making requests will use :tl:`InvokeWithTakeoutRequest` to wrap which making requests will use :tl:`InvokeWithTakeoutRequest` to wrap
@ -135,11 +143,14 @@ class AccountMethods(UserMethods):
before calling the method again. before calling the method again.
There's also a `success` property available in the takeout proxy There's also a `success` property available in the takeout proxy
object, so you can set the boolean takeout result that will be sent object, so from the `with` body you can set the boolean result that
back to Telegram, from within the `with` body. You're also able to will be sent back to Telegram. But if it's left ``None`` as by
set it to ``None`` so takeout will not be finished and its ID will default, then the action is based on the `finalize` parameter. If
be preserved for future usage as `client.session.takeout_id it's ``True`` then the takeout will be finished, and if no exception
<telethon.sessions.abstract.Session.takeout_id>`. Default is ``True``. occurred during it, then ``True`` will be considered as a result.
Otherwise, the takeout will not be finished and its ID will be
preserved for future usage as `client.session.takeout_id
<telethon.sessions.abstract.Session.takeout_id>`.
Args: Args:
contacts (`bool`): contacts (`bool`):
@ -186,7 +197,7 @@ class AccountMethods(UserMethods):
else: else:
request = None request = None
return _TakeoutClient(self, request) return _TakeoutClient(finalize, self, request)
async def end_takeout(self, success): async def end_takeout(self, success):
""" """
@ -195,8 +206,9 @@ class AccountMethods(UserMethods):
Returns: Returns:
``True`` if the operation was successful, ``False`` otherwise. ``True`` if the operation was successful, ``False`` otherwise.
""" """
client = _TakeoutClient(self, None) try:
result = await client( async with _TakeoutClient(True, self, None) as takeout:
functions.account.FinishTakeoutSessionRequest(success)) takeout.success = success
self.session.takeout_id = None except ValueError:
return result return False
return True