Refactor plugin management a bit, etc fixes/improvements

disco/bot/bot.py

@@ -262,7 +262,7 @@ class Bot(object):
                 content = content.replace('@everyone', '', 1)
             else:
                 for role in mention_roles:
-                    content = content.replace(role.mention, '', 1)
+                    content = content.replace('<@{}>'.format(role), '', 1)
 
             content = content.lstrip()
 
@@ -356,7 +356,7 @@ class Bot(object):
 
                 self.last_message_cache[msg.channel_id] = (msg, triggered)
 
-    def add_plugin(self, cls, config=None):
+    def add_plugin(self, cls, config=None, ctx=None):
         """
         Adds and loads a plugin, based on its class.
 
@@ -377,7 +377,7 @@ class Bot(object):
                 config = self.load_plugin_config(cls)
 
         self.plugins[cls.__name__] = cls(self, config)
-        self.plugins[cls.__name__].load()
+        self.plugins[cls.__name__].load(ctx or {})
         self.recompute()
 
     def rmv_plugin(self, cls):
@@ -392,9 +392,11 @@ class Bot(object):
         if cls.__name__ not in self.plugins:
             raise Exception('Cannot remove non-existant plugin: {}'.format(cls.__name__))
 
-        self.plugins[cls.__name__].unload()
+        ctx = {}
+        self.plugins[cls.__name__].unload(ctx)
         del self.plugins[cls.__name__]
         self.recompute()
+        return ctx
 
     def reload_plugin(self, cls):
         """
@@ -402,9 +404,9 @@ class Bot(object):
         """
         config = self.plugins[cls.__name__].config
 
-        self.rmv_plugin(cls)
+        ctx = self.rmv_plugin(cls)
         module = reload_module(inspect.getmodule(cls))
-        self.add_plugin(getattr(module, cls.__name__), config)
+        self.add_plugin(getattr(module, cls.__name__), config, ctx)
 
     def run_forever(self):
         """

disco/bot/plugin.py

@@ -1,4 +1,5 @@
 import six
+import types
 import gevent
 import inspect
 import weakref
@@ -18,8 +19,8 @@ class PluginDeco(object):
     Prio = Priority
 
     # TODO: dont smash class methods
-    @staticmethod
-    def add_meta_deco(meta):
+    @classmethod
+    def add_meta_deco(cls, meta):
         def deco(f):
             if not hasattr(f, 'meta'):
                 f.meta = []
@@ -153,6 +154,20 @@ class Plugin(LoggingClass, PluginDeco):
         self.storage = bot.storage
         self.config = config
 
+        # This is an array of all meta functions we sniff at init
+        self.meta_funcs = []
+
+        for name, member in inspect.getmembers(self, predicate=inspect.ismethod):
+            if hasattr(member, 'meta'):
+                self.meta_funcs.append(member)
+
+                # Unsmash local functions
+                if hasattr(Plugin, name):
+                    method = types.MethodType(getattr(Plugin, name), self, self.__class__)
+                    setattr(self, name, method)
+
+        self.bind_all()
+
     @property
     def name(self):
         return self.__class__.__name__
@@ -166,23 +181,21 @@ class Plugin(LoggingClass, PluginDeco):
         self._pre = {'command': [], 'listener': []}
         self._post = {'command': [], 'listener': []}
 
-        # TODO: when handling events/commands we need to track the greenlet in
-        #  the greenlets set so we can termiante long running commands/listeners
-        #  on reload.
-
-        for name, member in inspect.getmembers(self, predicate=inspect.ismethod):
-            if hasattr(member, 'meta'):
-                for meta in member.meta:
-                    if meta['type'] == 'listener':
-                        self.register_listener(member, meta['what'], meta['desc'], meta['priority'])
-                    elif meta['type'] == 'command':
-                        meta['kwargs']['update'] = True
-                        self.register_command(member, *meta['args'], **meta['kwargs'])
-                    elif meta['type'] == 'schedule':
-                        self.register_schedule(member, *meta['args'], **meta['kwargs'])
-                    elif meta['type'].startswith('pre_') or meta['type'].startswith('post_'):
-                        when, typ = meta['type'].split('_', 1)
-                        self.register_trigger(typ, when, member)
+        for member in self.meta_funcs:
+            for meta in member.meta:
+                self.bind_meta(member, meta)
+
+    def bind_meta(self, member, meta):
+        if meta['type'] == 'listener':
+            self.register_listener(member, meta['what'], meta['desc'], meta['priority'])
+        elif meta['type'] == 'command':
+            meta['kwargs']['update'] = True
+            self.register_command(member, *meta['args'], **meta['kwargs'])
+        elif meta['type'] == 'schedule':
+            self.register_schedule(member, *meta['args'], **meta['kwargs'])
+        elif meta['type'].startswith('pre_') or meta['type'].startswith('post_'):
+            when, typ = meta['type'].split('_', 1)
+            self.register_trigger(typ, when, member)
 
     def spawn(self, method, *args, **kwargs):
         obj = gevent.spawn(method, *args, **kwargs)
@@ -208,6 +221,7 @@ class Plugin(LoggingClass, PluginDeco):
         getattr(self, '_' + when)[typ].append(func)
 
     def _dispatch(self, typ, func, event, *args, **kwargs):
+        self.greenlets.add(gevent.getcurrent())
         self.ctx['plugin'] = self
 
         if hasattr(event, 'guild'):
@@ -302,13 +316,13 @@ class Plugin(LoggingClass, PluginDeco):
 
         self.schedules[func.__name__] = self.spawn(repeat)
 
-    def load(self):
+    def load(self, ctx):
         """
         Called when the plugin is loaded
         """
-        self.bind_all()
+        pass
 
-    def unload(self):
+    def unload(self, ctx):
         """
         Called when the plugin is unloaded
         """

disco/cli.py

@@ -42,6 +42,7 @@ def disco_main(run=False):
     from disco.bot import Bot, BotConfig
     from disco.gateway.sharder import AutoSharder
     from disco.util.token import is_valid_token
+    from holster.log import set_logging_levels
 
     if os.path.exists(args.config):
         config = ClientConfig.from_file(args.config)
@@ -61,6 +62,7 @@ def disco_main(run=False):
         return
 
     logging.basicConfig(level=logging.INFO)
+    set_logging_levels()
 
     client = Client(config)
 

disco/gateway/ipc/gipc.py

@@ -2,8 +2,7 @@ import random
 import gevent
 import string
 import weakref
-import marshal
-import types
+import dill
 
 from holster.enum import Enum
 
@@ -50,10 +49,10 @@ class GIPCProxy(LoggingClass):
             self.send(IPCMessageType.RESPONSE, (nonce, self.resolve(path)))
         elif mtype == IPCMessageType.EXECUTE:
             nonce, raw = data
-            func = types.FunctionType(marshal.loads(raw), globals(), nonce)
+            func = dill.loads(raw)
             try:
                 result = func(self.obj)
-            except Exception as e:
+            except Exception:
                 self.log.exception('Failed to EXECUTE: ')
                 result = None
 
@@ -74,7 +73,7 @@ class GIPCProxy(LoggingClass):
 
     def execute(self, func):
         nonce = get_random_str(32)
-        raw = marshal.dumps(func.func_code)
+        raw = dill.dumps(func)
         self.results[nonce] = result = gevent.event.AsyncResult()
         self.pipe.put((IPCMessageType.EXECUTE.value, (nonce, raw)))
         return result

disco/gateway/sharder.py

@@ -2,8 +2,10 @@ from __future__ import absolute_import
 
 import gipc
 import gevent
-import types
-import marshal
+import logging
+import dill
+
+from holster.log import set_logging_levels
 
 from disco.client import Client
 from disco.bot import Bot, BotConfig
@@ -13,7 +15,7 @@ from disco.gateway.ipc.gipc import GIPCProxy
 
 def run_on(id, proxy):
     def f(func):
-        return proxy.call(('run_on', ), id, marshal.dumps(func.func_code))
+        return proxy.call(('run_on', ), id, dill.dumps(func))
     return f
 
 
@@ -26,11 +28,11 @@ def run_self(bot):
 
 
 def run_shard(config, id, pipe):
-    import logging
     logging.basicConfig(
         level=logging.INFO,
         format='{} [%(levelname)s] %(asctime)s - %(name)s:%(lineno)d - %(message)s'.format(id)
     )
+    set_logging_levels()
 
     config.shard_id = id
     client = Client(config)
@@ -50,11 +52,11 @@ class AutoSharder(object):
         self.client = APIClient(config.token)
         self.shards = {}
         self.config.shard_count = self.client.gateway_bot_get()['shards']
-        self.config.shard_count = 10
-        self.test = 1
+        if self.config.shard_count > 1:
+            self.config.shard_count = 10
 
-    def run_on(self, id, funccode):
-        func = types.FunctionType(marshal.loads(funccode), globals(), '_run_on_temp')
+    def run_on(self, id, raw):
+        func = dill.loads(raw)
         return self.shards[id].execute(func).wait(timeout=15)
 
     def run(self):
@@ -65,7 +67,12 @@ class AutoSharder(object):
             self.start_shard(shard)
             gevent.sleep(6)
 
+        logging.basicConfig(
+            level=logging.INFO,
+            format='{} [%(levelname)s] %(asctime)s - %(name)s:%(lineno)d - %(message)s'.format(id)
+        )
+
     def start_shard(self, id):
-        cpipe, ppipe = gipc.pipe(duplex=True)
+        cpipe, ppipe = gipc.pipe(duplex=True, encoder=dill.dumps, decoder=dill.loads)
         gipc.start_process(run_shard, (self.config, id, cpipe))
         self.shards[id] = GIPCProxy(self, ppipe)

disco/state.py

@@ -243,6 +243,9 @@ class State(object):
         if event.member.guild_id not in self.guilds:
             return
 
+        if event.member.id not in self.guilds[event.member.guild_id].members:
+            return
+
         self.guilds[event.member.guild_id].members[event.member.id].update(event.member)
 
     def on_guild_member_remove(self, event):