From df932a21a352f7f0b1e0d62b258d6c7245cb486c Mon Sep 17 00:00:00 2001 From: Joseph Montanaro Date: Sat, 9 May 2020 10:45:07 -0700 Subject: [PATCH] streamline cleanup handling --- bot.py | 22 ++++++++++------------ 1 file changed, 10 insertions(+), 12 deletions(-) diff --git a/bot.py b/bot.py index 284ab52..db07c40 100644 --- a/bot.py +++ b/bot.py @@ -1,3 +1,4 @@ +import inspect import logging import os @@ -8,21 +9,18 @@ import heavynode class CustomBot(commands.Bot): def __init__(self, *args, **kwargs): - self.cleanup_sync = [] - self.cleanup_async = [] + self.cleanup = [] super().__init__(*args, **kwargs) - def add_cleanup_sync(self, func): - self.cleanup_sync.append(func) - - def add_cleanup_async(self, coro): - self.cleanup_async.append(coro) + def add_cleanup(self, callback): + self.cleanup.append(callback) async def close(self): - for func in self.cleanup_sync: - func() - for coro in self.cleanup_async: - await coro() + for callback in self.cleanup: + r = callback() + # coroutines etc. return an awaitable object + if inspect.isawaitable(r): + await r await super().close() @@ -32,7 +30,7 @@ server_id = 'd030737c' hn = heavynode.Client(os.environ['heavynode_token']) bot = CustomBot(command_prefix='!') -bot.add_cleanup_async(hn.shutdown) +bot.add_cleanup(hn.shutdown) @bot.command()