Update takeout to use less hacks

This commit is contained in:
Lonami Exo 2022-01-09 14:41:10 +01:00
parent 2db0725b98
commit be0da9b183
5 changed files with 118 additions and 136 deletions

View File

@ -697,6 +697,24 @@ If you were relying on any of the individual mixins that made up the client, suc
There is a single ``TelegramClient`` class now, containing everything you need.
The takeout context-manager has changed
---------------------------------------
It no longer has a finalize. All the requests made by the client in the same task will be wrapped,
not only those made through the proxy client returned by the context-manager.
This cleans up the (rather hacky) implementation, making use of Python's ``contextvar``. If you
still need the takeout session to persist, you should manually use the ``begin_takeout`` and
``end_takeout`` method.
If you want to ignore the currently-active takeout session in a task, toggle the following context
variable:
.. code-block:: python
telethon.ignore_takeout.set(True)
CdnDecrypter has been removed
-----------------------------

View File

@ -5,6 +5,7 @@ from ._misc import utils as _ # depends on helpers and _tl
from ._misc import hints as _ # depends on types/custom
from ._client.telegramclient import TelegramClient
from ._client.account import ignore_takeout
from . import version, events, errors, enums
__version__ = version.__version__

View File

@ -1,6 +1,7 @@
import functools
import inspect
import typing
from contextvars import ContextVar
from .users import _NOT_A_REQUEST
from .._misc import helpers, utils
@ -10,112 +11,43 @@ if typing.TYPE_CHECKING:
from .telegramclient import TelegramClient
ignore_takeout = ContextVar('ignore_takeout', default=False)
# TODO Make use of :tl:`InvokeWithMessagesRange` somehow
# For that, we need to use :tl:`GetSplitRanges` first.
class _TakeoutClient:
"""
Proxy object over the client.
"""
__PROXY_INTERFACE = ('__enter__', '__exit__', '__aenter__', '__aexit__')
def __init__(self, finalize, client, request):
# We use the name mangling for attributes to make them inaccessible
# from within the shadowed client object and to distinguish them from
# its own attributes where needed.
self.__finalize = finalize
self.__client = client
self.__request = request
self.__success = None
@property
def success(self):
return self.__success
@success.setter
def success(self, value):
self.__success = value
class _Takeout:
def __init__(self, client, kwargs):
self._client = client
self._kwargs = kwargs
async def __aenter__(self):
# Enter/Exit behaviour is "overrode", we don't want to call start.
client = self.__client
if client.session.takeout_id is None:
client.session.takeout_id = (await client(self.__request)).id
elif self.__request is not None:
raise ValueError("Can't send a takeout request while another "
"takeout for the current session still not been finished yet.")
return self
await self._client.begin_takeout(**kwargs)
return self._client
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:
result = await self(_tl.fn.account.FinishTakeoutSession(
self.__success))
if not result:
raise ValueError("Failed to finish the takeout.")
self.session.takeout_id = None
async def __call__(self, request, ordered=False):
takeout_id = self.__client.session.takeout_id
if takeout_id is None:
raise ValueError('Takeout mode has not been initialized '
'(are you calling outside of "with"?)')
single = not utils.is_list_like(request)
requests = ((request,) if single else request)
wrapped = []
for r in requests:
if not isinstance(r, _tl.TLRequest):
raise _NOT_A_REQUEST()
await r.resolve(self, utils)
wrapped.append(_tl.fn.InvokeWithTakeout(takeout_id, r))
return await self.__client(
wrapped[0] if single else wrapped, ordered=ordered)
def __getattribute__(self, name):
# We access class via type() because __class__ will recurse infinitely.
# Also note that since we've name-mangled our own class attributes,
# they'll be passed to __getattribute__() as already decorated. For
# example, 'self.__client' will be passed as '_TakeoutClient__client'.
# https://docs.python.org/3/tutorial/classes.html#private-variables
if name.startswith('__') and name not in type(self).__PROXY_INTERFACE:
raise AttributeError # force call of __getattr__
# Try to access attribute in the proxy object and check for the same
# attribute in the shadowed object (through our __getattr__) if failed.
return super().__getattribute__(name)
def __getattr__(self, name):
value = getattr(self.__client, name)
if inspect.ismethod(value):
# Emulate bound methods behavior by partially applying our proxy
# class as the self parameter instead of the client.
return functools.partial(
getattr(self.__client.__class__, name), self)
return value
def __setattr__(self, name, value):
if name.startswith('_{}__'.format(type(self).__name__.lstrip('_'))):
# This is our own name-mangled attribute, keep calm.
return super().__setattr__(name, value)
return setattr(self.__client, name, value)
await self._client.end_takeout(success=exc_type is None)
def takeout(
self: 'TelegramClient',
finalize: bool = True,
*,
contacts: bool = None,
users: bool = None,
chats: bool = None,
megagroups: bool = None,
channels: bool = None,
files: bool = None,
max_file_size: bool = None) -> 'TelegramClient':
request_kwargs = dict(
def takeout(self: 'TelegramClient', **kwargs):
return _Takeout(self, kwargs)
async def begin_takeout(
self: 'TelegramClient',
*,
contacts: bool = None,
users: bool = None,
chats: bool = None,
megagroups: bool = None,
channels: bool = None,
files: bool = None,
max_file_size: bool = None,
) -> 'TelegramClient':
if takeout_active():
raise ValueError('a previous takeout session was already active')
self._session_state.takeout_id = (await client(
contacts=contacts,
message_users=users,
message_chats=chats,
@ -123,21 +55,19 @@ def takeout(
message_channels=channels,
files=files,
file_max_size=max_file_size
)
arg_specified = (arg is not None for arg in request_kwargs.values())
)).id
if self.session.takeout_id is None or any(arg_specified):
request = _tl.fn.account.InitTakeoutSession(
**request_kwargs)
else:
request = None
return _TakeoutClient(finalize, self, request)
def takeout_active(self: 'TelegramClient') -> bool:
return self._session_state.takeout_id is not None
async def end_takeout(self: 'TelegramClient', success: bool) -> bool:
try:
async with _TakeoutClient(True, self, None) as takeout:
takeout.success = success
except ValueError:
return False
return True
if not takeout_active():
raise ValueError('no previous takeout session was active')
result = await self(_tl.fn.account.FinishTakeoutSession(success))
if not result:
raise ValueError("could not end the active takeout session")
self._session_state.takeout_id = None

View File

@ -174,7 +174,6 @@ class TelegramClient:
@forward_call(account.takeout)
def takeout(
self: 'TelegramClient',
finalize: bool = True,
*,
contacts: bool = None,
users: bool = None,
@ -184,14 +183,39 @@ class TelegramClient:
files: bool = None,
max_file_size: bool = None) -> 'TelegramClient':
"""
Returns a :ref:`telethon-client` which calls methods behind a takeout session.
Returns a context-manager which calls `TelegramClient.begin_takeout`
on enter and `TelegramClient.end_takeout` on exit. The same errors
and conditions apply.
It does so by creating a proxy object over the current client through
which making requests will use :tl:`InvokeWithTakeout` to wrap
them. In other words, returns the current client modified so that
requests are done as a takeout:
This is useful for the common case of not wanting the takeout to
persist (although it still might if a disconnection occurs before it
can be ended).
Some of the calls made through the takeout session will have lower
Example
.. code-block:: python
async with client.takeout():
async for message in client.iter_messages(chat, wait_time=0):
... # Do something with the message
"""
@forward_call(account.begin_takeout)
def begin_takeout(
self: 'TelegramClient',
*,
contacts: bool = None,
users: bool = None,
chats: bool = None,
megagroups: bool = None,
channels: bool = None,
files: bool = None,
max_file_size: bool = None) -> 'TelegramClient':
"""
Begin a takeout session. All subsequent requests made by the client
will be behind a takeout session. The takeout session will persist
in the session file, until `TelegramClient.end_takeout` is used.
When the takeout session is enabled, some requests will have lower
flood limits. This is useful if you want to export the data from
conversations or mass-download media, since the rate limits will
be lower. Only some requests will be affected, and you will need
@ -206,20 +230,16 @@ class TelegramClient:
can then access ``e.seconds`` to know how long you should wait for
before calling the method again.
There's also a `success` property available in the takeout proxy
object, so from the `with` body you can set the boolean result that
will be sent back to Telegram. But if it's left `None` as by
default, then the action is based on the `finalize` parameter. If
it's `True` then the takeout will be finished, and if no exception
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 in the session.
If you want to ignore the currently-active takeout session in a task,
toggle the following context variable:
.. code-block:: python
telethon.ignore_takeout.set(True)
An error occurs if ``TelegramClient.takeout_active`` was already ``True``.
Arguments
finalize (`bool`):
Whether the takeout session should be finalized upon
exit or not.
contacts (`bool`):
Set to `True` if you plan on downloading contacts.
@ -253,17 +273,26 @@ class TelegramClient:
from telethon import errors
try:
async with client.takeout() as takeout:
await client.get_messages('me') # normal call
await takeout.get_messages('me') # wrapped through takeout (less limits)
await client.begin_takeout()
async for message in takeout.iter_messages(chat, wait_time=0):
... # Do something with the message
await client.get_messages('me') # wrapped through takeout (less limits)
async for message in client.iter_messages(chat, wait_time=0):
... # Do something with the message
await client.end_takeout(success=True)
except errors.TakeoutInitDelayError as e:
print('Must wait', e.seconds, 'before takeout')
except Exception:
await client.end_takeout(success=False)
"""
@property
def takeout_active(self: 'TelegramClient') -> bool:
return account.takeout_active(self)
@forward_call(account.end_takeout)
async def end_takeout(self: 'TelegramClient', success: bool) -> bool:
"""

View File

@ -9,6 +9,7 @@ from ..errors._rpcbase import RpcError, ServerError, FloodError, InvalidDcError,
from .._misc import helpers, utils, hints
from .._sessions.types import Entity
from .. import errors, _tl
from .account import ignore_takeout
_NOT_A_REQUEST = lambda: TypeError('You can only invoke requests, not types!')
@ -52,6 +53,9 @@ async def _call(self: 'TelegramClient', sender, request, ordered=False, flood_sl
else:
raise errors.FLOOD_WAIT(420, f'FLOOD_WAIT_{diff}', request=r)
if self._session_state.takeout_id and not ignore_takeout.get():
r = _tl.fn.InvokeWithTakeout(self._session_state.takeout_id, r)
if self._no_updates:
r = _tl.fn.InvokeWithoutUpdates(r)