From 6d26d73e24bb70493546948594ff848d0b4d8680 Mon Sep 17 00:00:00 2001 From: Rapptz Date: Mon, 14 Dec 2015 16:26:15 -0500 Subject: [PATCH] Client.run now does cleanup when KeyboardInterrupt is encountered. --- discord/client.py | 64 +++++++++++++++++++++++++++++------------ discord/voice_client.py | 21 ++++++++------ 2 files changed, 58 insertions(+), 27 deletions(-) diff --git a/discord/client.py b/discord/client.py index af7b347b7..90fd5f1d3 100644 --- a/discord/client.py +++ b/discord/client.py @@ -262,8 +262,13 @@ class Client: def _run_event(self, event, *args, **kwargs): try: yield from getattr(self, event)(*args, **kwargs) - except Exception as e: - yield from self.on_error(event, *args, **kwargs) + except asyncio.CancelledError: + pass + except Exception: + try: + yield from self.on_error(event, *args, **kwargs) + except asyncio.CancelledError: + pass def dispatch(self, event, *args, **kwargs): log.debug('Dispatching event {}'.format(event)) @@ -278,16 +283,19 @@ class Client: @asyncio.coroutine def keep_alive_handler(self, interval): - while not self._closed: - payload = { - 'op': 1, - 'd': int(time.time()) - } - - msg = 'Keeping websocket alive with timestamp {}' - log.debug(msg.format(payload['d'])) - yield from self.ws.send(utils.to_json(payload)) - yield from asyncio.sleep(interval) + try: + while not self._closed: + payload = { + 'op': 1, + 'd': int(time.time()) + } + + msg = 'Keeping websocket alive with timestamp {}' + log.debug(msg.format(payload['d'])) + yield from self.ws.send(utils.to_json(payload)) + yield from asyncio.sleep(interval) + except asyncio.CancelledError: + pass @asyncio.coroutine def on_error(self, event_method, *args, **kwargs): @@ -662,14 +670,34 @@ class Client: """A blocking call that abstracts away the `event loop`_ initialisation from you. - Equivalent to: :: + If you want more control over the event loop then this + function should not be used. Use :meth:`start` coroutine + or :meth:`connect` + :meth:`login`. - loop.run_until_complete(start(email, password)) - loop.close() - """ + Roughly Equivalent to: :: - self.loop.run_until_complete(self.start(email, password)) - self.loop.close() + try: + loop.run_until_complete(start(email, password)) + except KeyboardInterrupt: + loop.run_until_complete(logout()) + # cancel all tasks lingering + finally: + loop.close() + """ + try: + self.loop.run_until_complete(self.start(email, password)) + except KeyboardInterrupt: + self.loop.run_until_complete(self.logout()) + pending = asyncio.Task.all_tasks() + gathered = asyncio.gather(*pending) + try: + gathered.cancel() + self.loop.run_forever() + gathered.exception() + except: + pass + finally: + self.loop.close() # event registration diff --git a/discord/voice_client.py b/discord/voice_client.py index e0e0a85ba..22471c997 100644 --- a/discord/voice_client.py +++ b/discord/voice_client.py @@ -175,16 +175,19 @@ class VoiceClient: @asyncio.coroutine def keep_alive_handler(self, delay): - while True: - payload = { - 'op': 3, - 'd': int(time.time()) - } + try: + while True: + payload = { + 'op': 3, + 'd': int(time.time()) + } - msg = 'Keeping voice websocket alive with timestamp {}' - log.debug(msg.format(payload['d'])) - yield from self.ws.send(utils.to_json(payload)) - yield from asyncio.sleep(delay) + msg = 'Keeping voice websocket alive with timestamp {}' + log.debug(msg.format(payload['d'])) + yield from self.ws.send(utils.to_json(payload)) + yield from asyncio.sleep(delay) + except asyncio.CancelledError: + pass @asyncio.coroutine def received_message(self, msg):