From 2d275989cb372b66b7e0e2f20a7262585a623e90 Mon Sep 17 00:00:00 2001 From: Lonami Exo Date: Sat, 29 Sep 2018 13:36:05 +0200 Subject: [PATCH] Properly handle cancellation in _ReadyQueue --- telethon/helpers.py | 16 ++++++++++------ telethon/network/connection/connection.py | 4 ++-- 2 files changed, 12 insertions(+), 8 deletions(-) diff --git a/telethon/helpers.py b/telethon/helpers.py index e0cced62..861e56ea 100644 --- a/telethon/helpers.py +++ b/telethon/helpers.py @@ -111,12 +111,16 @@ class _ReadyQueue: Returns a list of all the items added to the queue until now and clears the list from the queue itself. Returns ``None`` if cancelled. """ - ready = asyncio.ensure_future(self._ready.wait(), loop=self._loop) - done, pending = await asyncio.wait( - [ready, cancellation], - return_when=asyncio.FIRST_COMPLETED, - loop=self._loop - ) + ready = self._loop.create_task(self._ready.wait()) + try: + done, pending = await asyncio.wait( + [ready, cancellation], + return_when=asyncio.FIRST_COMPLETED, + loop=self._loop + ) + except asyncio.CancelledError: + done = [cancellation] + if cancellation in done: ready.cancel() return None diff --git a/telethon/network/connection/connection.py b/telethon/network/connection/connection.py index abe0fa60..1dfd1698 100644 --- a/telethon/network/connection/connection.py +++ b/telethon/network/connection/connection.py @@ -51,8 +51,8 @@ class Connection(abc.ABC): @property def disconnected(self): if not self._disconnected_future: - self._disconnected_future = asyncio.ensure_future( - self._disconnected.wait(), loop=self._loop) + self._disconnected_future = \ + self._loop.create_task(self._disconnected.wait()) return self._disconnected_future def clone(self):