streamline cleanup handling
This commit is contained in:
parent
485aae06c9
commit
df932a21a3
22
bot.py
22
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()
|
||||
|
Loading…
x
Reference in New Issue
Block a user