streamline cleanup handling

This commit is contained in:
Joseph Montanaro 2020-05-09 10:45:07 -07:00
parent 485aae06c9
commit df932a21a3

22
bot.py
View File

@ -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()