From 9090ede5dbb68f22947ca26af618c9278a7df28f Mon Sep 17 00:00:00 2001 From: Lonami Exo Date: Sat, 13 Apr 2019 10:53:33 +0200 Subject: [PATCH] Reduce __enter__/__exit__ boilerplate for sync ctx managers --- telethon/client/account.py | 17 ++++------------- telethon/client/auth.py | 16 +++------------- telethon/client/chats.py | 15 +++------------ telethon/helpers.py | 19 +++++++++++++++++++ telethon/tl/custom/conversation.py | 17 ++++------------- 5 files changed, 33 insertions(+), 51 deletions(-) diff --git a/telethon/client/account.py b/telethon/client/account.py index 703e96f0..69ced957 100644 --- a/telethon/client/account.py +++ b/telethon/client/account.py @@ -2,7 +2,7 @@ import functools import inspect from .users import UserMethods, _NOT_A_REQUEST -from .. import utils +from .. import helpers, utils from ..tl import functions, TLRequest @@ -31,15 +31,6 @@ class _TakeoutClient: def success(self, value): self.__success = value - def __enter__(self): - if self.__client.loop.is_running(): - raise RuntimeError( - 'You must use "async with" if the event loop ' - 'is running (i.e. you are inside an "async def")' - ) - - return self.__client.loop.run_until_complete(self.__aenter__()) - async def __aenter__(self): # Enter/Exit behaviour is "overrode", we don't want to call start. client = self.__client @@ -50,9 +41,6 @@ class _TakeoutClient: "takeout for the current session still not been finished yet.") return self - def __exit__(self, *args): - return self.__client.loop.run_until_complete(self.__aexit__(*args)) - async def __aexit__(self, exc_type, exc_value, traceback): if self.__success is None and self.__finalize: self.__success = exc_type is None @@ -64,6 +52,9 @@ class _TakeoutClient: raise ValueError("Failed to finish the takeout.") self.session.takeout_id = None + __enter__ = helpers._sync_enter + __exit__ = helpers._sync_exit + async def __call__(self, request, ordered=False): takeout_id = self.__client.session.takeout_id if takeout_id is None: diff --git a/telethon/client/auth.py b/telethon/client/auth.py index a035833a..d30ba6d2 100644 --- a/telethon/client/auth.py +++ b/telethon/client/auth.py @@ -536,23 +536,13 @@ class AuthMethods(MessageParseMethods, UserMethods): # region with blocks - def __enter__(self): - if self._loop.is_running(): - raise RuntimeError( - 'You must use "async with" if the event loop ' - 'is running (i.e. you are inside an "async def")' - ) - - return self.start() - async def __aenter__(self): return await self.start() - def __exit__(self, *args): - # No loop.run_until_complete; it's already syncified - self.disconnect() - async def __aexit__(self, *args): await self.disconnect() + __enter__ = helpers._sync_enter + __exit__ = helpers._sync_exit + # endregion diff --git a/telethon/client/chats.py b/telethon/client/chats.py index 24bc4100..b929ed47 100644 --- a/telethon/client/chats.py +++ b/telethon/client/chats.py @@ -3,7 +3,7 @@ import itertools import string from .users import UserMethods -from .. import utils +from .. import helpers, utils from ..requestiter import RequestIter from ..tl import types, functions, custom @@ -69,17 +69,8 @@ class _ChatAction: self._task = None - def __enter__(self): - if self._client.loop.is_running(): - raise RuntimeError( - 'You must use "async with" if the event loop ' - 'is running (i.e. you are inside an "async def")' - ) - - return self._client.loop.run_until_complete(self.__aenter__()) - - def __exit__(self, *args): - return self._client.loop.run_until_complete(self.__aexit__(*args)) + __enter__ = helpers._sync_enter + __exit__ = helpers._sync_exit async def _update(self): try: diff --git a/telethon/helpers.py b/telethon/helpers.py index 33276b8f..531cc421 100644 --- a/telethon/helpers.py +++ b/telethon/helpers.py @@ -102,6 +102,25 @@ async def _cancel(log, **tasks): except Exception: log.exception('Unhandled exception from %s after cancel', name) + +def _sync_enter(self): + """ + Helps to cut boilerplate on async context + managers that offer synchronous variants. + """ + if self._client.loop.is_running(): + raise RuntimeError( + 'You must use "async with" if the event loop ' + 'is running (i.e. you are inside an "async def")' + ) + + return self._client.loop.run_until_complete(self.__aenter__()) + + +def _sync_exit(self, *args): + return self._client.loop.run_until_complete(self.__aexit__(*args)) + + # endregion # region Cryptographic related utils diff --git a/telethon/tl/custom/conversation.py b/telethon/tl/custom/conversation.py index e3147803..d8d93b13 100644 --- a/telethon/tl/custom/conversation.py +++ b/telethon/tl/custom/conversation.py @@ -3,7 +3,7 @@ import itertools import time from .chatgetter import ChatGetter -from ... import utils, errors +from ... import helpers, utils, errors # Sometimes the edits arrive very fast (within the same second). @@ -403,15 +403,6 @@ class Conversation(ChatGetter): else: fut.cancel() - def __enter__(self): - if self._client.loop.is_running(): - raise RuntimeError( - 'You must use "async with" if the event loop ' - 'is running (i.e. you are inside an "async def")' - ) - - return self._client.loop.run_until_complete(self.__aenter__()) - async def __aenter__(self): self._input_chat = \ await self._client.get_input_entity(self._input_chat) @@ -447,9 +438,6 @@ class Conversation(ChatGetter): """Cancels the current conversation and exits the context manager.""" raise _ConversationCancelled() - def __exit__(self, *args): - return self._client.loop.run_until_complete(self.__aexit__(*args)) - async def __aexit__(self, exc_type, exc_val, exc_tb): chat_id = utils.get_peer_id(self._chat_peer) if self._client._ids_in_conversations[chat_id] == 1: @@ -461,6 +449,9 @@ class Conversation(ChatGetter): self._cancel_all() return isinstance(exc_val, _ConversationCancelled) + __enter__ = helpers._sync_enter + __exit__ = helpers._sync_exit + class _ConversationCancelled(InterruptedError): pass