Update takeout to use less hacks

readthedocs/misc/v2-migration-guide.rst

@@ -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
 -----------------------------
 

telethon/__init__.py

@@ -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__

telethon/_client/account.py

@@ -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:
+class _Takeout:
-    """
+    def __init__(self, client, kwargs):
-    Proxy object over the client.
+        self._client = client
-    """
+        self._kwargs = kwargs
-    __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
 
     async def __aenter__(self):
-        # Enter/Exit behaviour is "overrode", we don't want to call start.
+        await self._client.begin_takeout(**kwargs)
-        client = self.__client
+        return 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
 
     async def __aexit__(self, exc_type, exc_value, traceback):
-        if self.__success is None and self.__finalize:
+        await self._client.end_takeout(success=exc_type is None)
-            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)
 
 
-def takeout(
+def takeout(self: 'TelegramClient', **kwargs):
-        self: 'TelegramClient',
+    return _Takeout(self, kwargs)
-        finalize: bool = True,
+
-        *,
+
-        contacts: bool = None,
+async def begin_takeout(
-        users: bool = None,
+    self: 'TelegramClient',
-        chats: bool = None,
+    *,
-        megagroups: bool = None,
+    contacts: bool = None,
-        channels: bool = None,
+    users: bool = None,
-        files: bool = None,
+    chats: bool = None,
-        max_file_size: bool = None) -> 'TelegramClient':
+    megagroups: bool = None,
-    request_kwargs = dict(
+    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
-    )
+    )).id
-    arg_specified = (arg is not None for arg in request_kwargs.values())
 
-    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:
+    if not takeout_active():
-        async with _TakeoutClient(True, self, None) as takeout:
+        raise ValueError('no previous takeout session was active')
-            takeout.success = success
+
-    except ValueError:
+    result = await self(_tl.fn.account.FinishTakeoutSession(success))
-        return False
+    if not result:
-    return True
+        raise ValueError("could not end the active takeout session")
+
+    self._session_state.takeout_id = None

telethon/_client/telegramclient.py

@@ -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
+        This is useful for the common case of not wanting the takeout to
-        which making requests will use :tl:`InvokeWithTakeout` to wrap
+        persist (although it still might if a disconnection occurs before it
-        them. In other words, returns the current client modified so that
+        can be ended).
-        requests are done as a takeout:
 
-        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
+        If you want to ignore the currently-active takeout session in a task,
-        object, so from the `with` body you can set the boolean result that
+        toggle the following context variable:
-        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
+        .. code-block:: python
-        it's `True` then the takeout will be finished, and if no exception
+
-        occurred during it, then `True` will be considered as a result.
+            telethon.ignore_takeout.set(True)
-        Otherwise, the takeout will not be finished and its ID will be
+
-        preserved for future usage in the session.
+        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.begin_takeout()
-                        await client.get_messages('me')  # normal call
-                        await takeout.get_messages('me')  # wrapped through takeout (less limits)
 
-                        async for message in takeout.iter_messages(chat, wait_time=0):
+                    await client.get_messages('me')  # wrapped through takeout (less limits)
-                            ...  # Do something with the message
+
+                    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:
         """

telethon/_client/users.py

@@ -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)