streamline cleanup handling

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 = []
-        self.cleanup_async = []
         super().__init__(*args, **kwargs)
 
-    def add_cleanup_sync(self, func):
+    def add_cleanup(self, callback):
-        self.cleanup_sync.append(func)
+        self.cleanup.append(callback)
-
-    def add_cleanup_async(self, coro):
-        self.cleanup_async.append(coro)
 
     async def close(self):
-        for func in self.cleanup_sync:
+        for callback in self.cleanup:
-            func()
+            r = callback()
-        for coro in self.cleanup_async:
+            # coroutines etc. return an awaitable object
-            await coro()
+            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()