diff --git a/disco/bot/bot.py b/disco/bot/bot.py index 99ac0b5..70abe0b 100644 --- a/disco/bot/bot.py +++ b/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): """ diff --git a/disco/bot/plugin.py b/disco/bot/plugin.py index 5ee6803..8451b82 100644 --- a/disco/bot/plugin.py +++ b/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 """ diff --git a/disco/cli.py b/disco/cli.py index f2ef504..487de95 100644 --- a/disco/cli.py +++ b/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) diff --git a/disco/gateway/ipc/gipc.py b/disco/gateway/ipc/gipc.py index 4fdd49e..e23206b 100644 --- a/disco/gateway/ipc/gipc.py +++ b/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 diff --git a/disco/gateway/sharder.py b/disco/gateway/sharder.py index 987ca73..cb387a4 100644 --- a/disco/gateway/sharder.py +++ b/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) diff --git a/disco/state.py b/disco/state.py index ae50698..c0e5338 100644 --- a/disco/state.py +++ b/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):