Fix begin/end_takeout

This commit is contained in:
Lonami Exo 2022-02-18 19:21:52 +01:00
parent d6669a1172
commit d1dba60278
2 changed files with 9 additions and 7 deletions

View File

@ -2,6 +2,7 @@ import functools
import inspect import inspect
import typing import typing
import dataclasses import dataclasses
import asyncio
from contextvars import ContextVar from contextvars import ContextVar
from .._misc import helpers, utils from .._misc import helpers, utils
@ -22,7 +23,7 @@ class _Takeout:
self._kwargs = kwargs self._kwargs = kwargs
async def __aenter__(self): async def __aenter__(self):
await self._client.begin_takeout(**kwargs) await self._client.begin_takeout(**self._kwargs)
return self._client return self._client
async def __aexit__(self, exc_type, exc_value, traceback): async def __aexit__(self, exc_type, exc_value, traceback):
@ -44,10 +45,10 @@ async def begin_takeout(
files: bool = None, files: bool = None,
max_file_size: bool = None, max_file_size: bool = None,
) -> 'TelegramClient': ) -> 'TelegramClient':
if takeout_active(): if self.takeout_active:
raise ValueError('a previous takeout session was already active') raise ValueError('a previous takeout session was already active')
await self._replace_session_state(takeout_id=(await client( takeout = await self(_tl.fn.account.InitTakeoutSession(
contacts=contacts, contacts=contacts,
message_users=users, message_users=users,
message_chats=chats, message_chats=chats,
@ -55,15 +56,16 @@ async def begin_takeout(
message_channels=channels, message_channels=channels,
files=files, files=files,
file_max_size=max_file_size file_max_size=max_file_size
)).id) ))
await self._replace_session_state(takeout_id=takeout.id)
def takeout_active(self: 'TelegramClient') -> bool: def takeout_active(self: 'TelegramClient') -> bool:
return self._session_state.takeout_id is not None return self._session_state.takeout_id is not None
async def end_takeout(self: 'TelegramClient', success: bool) -> bool: async def end_takeout(self: 'TelegramClient', *, success: bool) -> bool:
if not takeout_active(): if not self.takeout_active:
raise ValueError('no previous takeout session was active') raise ValueError('no previous takeout session was active')
result = await self(_tl.fn.account.FinishTakeoutSession(success)) result = await self(_tl.fn.account.FinishTakeoutSession(success))

View File

@ -289,7 +289,7 @@ class TelegramClient:
return account.takeout_active(self) return account.takeout_active(self)
@forward_call(account.end_takeout) @forward_call(account.end_takeout)
async def end_takeout(self: 'TelegramClient', success: bool) -> bool: async def end_takeout(self: 'TelegramClient', *, success: bool) -> bool:
""" """
Finishes the current takeout session. Finishes the current takeout session.