Better autosharding interface, more fixes

disco/bot/bot.py

@@ -113,6 +113,7 @@ class Bot(object):
     def __init__(self, client, config=None):
         self.client = client
         self.config = config or BotConfig()
+        self.shards = {}
 
         # The context carries information about events in a threadlocal storage
         self.ctx = ThreadLocal()

disco/bot/command.py

@@ -6,7 +6,7 @@ from disco.bot.parser import ArgumentSet, ArgumentError
 from disco.util.functional import cached_property
 
 REGEX_FMT = '({})'
-ARGS_REGEX = '( (.*)$|$)'
+ARGS_REGEX = '( ((?:\n|.)*)$|$)'
 MENTION_RE = re.compile('<@!?([0-9]+)>')
 
 CommandLevels = Enum(
@@ -109,7 +109,7 @@ class Command(object):
         self.triggers = [trigger]
         self.update(*args, **kwargs)
 
-    def update(self, args=None, level=None, aliases=None, group=None, is_regex=None):
+    def update(self, args=None, level=None, aliases=None, group=None, is_regex=None, oob=False):
         self.triggers += aliases or []
 
         def resolve_role(ctx, id):
@@ -127,6 +127,7 @@ class Command(object):
         self.level = level
         self.group = group
         self.is_regex = is_regex
+        self.oob = oob
 
     @staticmethod
     def mention_type(getters, force=False):

disco/bot/plugin.py

@@ -206,6 +206,8 @@ class Plugin(LoggingClass, PluginDeco):
         """
         Executes a CommandEvent this plugin owns
         """
+        if not event.command.oob:
+            self.greenlets.add(gevent.getcurrent())
         try:
             return event.command.execute(event)
         except CommandError as e:
@@ -221,7 +223,9 @@ class Plugin(LoggingClass, PluginDeco):
         getattr(self, '_' + when)[typ].append(func)
 
     def _dispatch(self, typ, func, event, *args, **kwargs):
-        self.greenlets.add(gevent.getcurrent())
+        # TODO: this is ugly
+        if typ != 'command':
+            self.greenlets.add(gevent.getcurrent())
         self.ctx['plugin'] = self
 
         if hasattr(event, 'guild'):

disco/gateway/sharder.py

@@ -1,9 +1,11 @@
 from __future__ import absolute_import
 
+import six
 import gipc
 import gevent
 import logging
 import dill
+import types
 
 from holster.log import set_logging_levels
 
@@ -11,11 +13,34 @@ from disco.client import Client
 from disco.bot import Bot, BotConfig
 from disco.api.client import APIClient
 from disco.gateway.ipc.gipc import GIPCProxy
+from disco.util.snowflake import calculate_shard
+
+
+def dump_function(func):
+    if six.PY3:
+        return dill.dumps((
+            func.__code__,
+            func.__name__,
+            func.__defaults__,
+            func.__closure__,
+        ))
+    else:
+        return dill.dumps((
+            func.func_code,
+            func.func_name,
+            func.func_defaults,
+            func.func_closure
+        ))
+
+
+def load_function(func):
+    code, name, defaults, closure = dill.loads(func)
+    return types.FunctionType(code, globals(), name, defaults, closure)
 
 
 def run_on(id, proxy):
     def f(func):
-        return proxy.call(('run_on', ), id, dill.dumps(func))
+        return proxy.call(('run_on', ), id, dump_function(func))
     return f
 
 
@@ -38,14 +63,36 @@ def run_shard(config, id, pipe):
     client = Client(config)
     bot = Bot(client, BotConfig(config.bot))
     bot.sharder = GIPCProxy(bot, pipe)
-    bot.shards = {
-        i: run_on(i, bot.sharder) for i in range(config.shard_count)
-        if i != id
-    }
-    bot.shards[id] = run_self(bot)
+    bot.shards = ShardHelper(config.shard_count, bot)
     bot.run_forever()
 
 
+class ShardHelper(object):
+    def __init__(self, count, bot):
+        self.count = count
+        self.bot = bot
+
+    def keys(self):
+        for id in xrange(self.count):
+            yield id
+
+    def on(self, id, func):
+        if id == self.bot.client.config.shard_id:
+            result = gevent.event.AsyncResult()
+            result.set(func(self.bot))
+            return result
+
+        return self.bot.sharder.call(('run_on', ), id, dump_function(func))
+
+    def all(self, func, timeout=None):
+        pool = gevent.pool.Pool(self.count)
+        return dict(zip(range(self.count), pool.imap(lambda i: self.on(i, func).wait(timeout=timeout), range(self.count))))
+
+    def for_id(self, id, func):
+        shard = calculate_shard(self.count, id)
+        return self.on(shard, func)
+
+
 class AutoSharder(object):
     def __init__(self, config):
         self.config = config
@@ -56,7 +103,8 @@ class AutoSharder(object):
             self.config.shard_count = 10
 
     def run_on(self, id, raw):
-        func = dill.loads(raw)
+        func = load_function(raw)
+        # func = dill.loads(raw)
         return self.shards[id].execute(func).wait(timeout=15)
 
     def run(self):

disco/types/guild.py

@@ -245,6 +245,7 @@ class Guild(SlottedModel, Permissible):
     roles = Field(dictof(Role, key='id'))
     emojis = Field(dictof(Emoji, key='id'))
     voice_states = Field(dictof(VoiceState, key='session_id'))
+    member_count = Field(int)
 
     synced = Field(bool, default=False)