Fixes, Cleanup, Plugin Storage

The biggest part of this commit is a plugin storage subsystem, which at
this point I'm fairly happy with. I've iterated on this a couple times,
and the final result has a very clean/simple interface, is easy to
extend to different data stores, and has a very few minimal number of
grokable edge cases.

- Storage subsytem
- Fix command group abbreviations
- Fix reconnecting in the GatewaySocket
- Add pickle support to serializer

.gitignore

@@ -2,3 +2,4 @@ build/
 dist/
 disco*.egg-info/
 docs/_build
+storage.db

disco/bot/bot.py

@@ -84,9 +84,8 @@ class BotConfig(Config):
     plugin_config_dir = 'config'
 
     storage_enabled = False
-    storage_backend = 'memory'
-    storage_autosave = True
-    storage_autosave_interval = 120
+    storage_provider = 'memory'
+    storage_config = {}
 
 
 class Bot(object):
@@ -184,9 +183,10 @@ def recompute(self):
         """
         Called when a plugin is loaded/unloaded to recompute internal state.
         """
-        self.compute_group_abbrev()
         if self.config.commands_group_abbrev:
-            self.compute_command_matches_re()
+            self.compute_group_abbrev()
+
+        self.compute_command_matches_re()
 
     def compute_group_abbrev(self):
         """

disco/bot/command.py

@@ -165,10 +165,10 @@ def regex(self):
         else:
             group = ''
             if self.group:
-                if self.group in self.plugin.bot.group_abbrev.get(self.group):
-                    group = '{}(?:\w+)? '.format(self.group)
+                if self.group in self.plugin.bot.group_abbrev:
+                    group = '{}(?:\w+)? '.format(self.plugin.bot.group_abbrev.get(self.group))
                 else:
-                    group = self.group
+                    group = self.group + ' '
             return REGEX_FMT.format('|'.join(['^' + group + trigger for trigger in self.triggers]) + ARGS_REGEX)
 
     def execute(self, event):

disco/bot/plugin.py

@@ -17,6 +17,7 @@ class PluginDeco(object):
     """
     Prio = Priority
 
+    # TODO: dont smash class methods
     @staticmethod
     def add_meta_deco(meta):
         def deco(f):
@@ -152,6 +153,10 @@ def __init__(self, bot, config):
         self.storage = bot.storage
         self.config = config
 
+    @property
+    def name(self):
+        return self.__class__.__name__
+
     def bind_all(self):
         self.listeners = []
         self.commands = {}
@@ -188,6 +193,7 @@ def execute(self, event):
         """
         Executes a CommandEvent this plugin owns
         """
+        self.ctx['plugin'] = self
         self.ctx['guild'] = event.guild
         self.ctx['channel'] = event.channel
         self.ctx['user'] = event.author

disco/bot/providers/__init__.py

@@ -0,0 +1,15 @@
+import inspect
+import importlib
+
+from .base import BaseProvider
+
+
+def load_provider(name):
+    try:
+        mod = importlib.import_module('disco.bot.providers.' + name)
+    except ImportError:
+        mod = importlib.import_module(name)
+
+    for entry in filter(inspect.isclass, map(lambda i: getattr(mod, i), dir(mod))):
+        if issubclass(entry, BaseProvider) and entry != BaseProvider:
+            return entry

disco/bot/providers/base.py

@@ -0,0 +1,136 @@
+import six
+import pickle
+
+from six.moves import map
+
+from UserDict import UserDict
+
+
+ROOT_SENTINEL = u'\u200B'
+SEP_SENTINEL = u'\u200D'
+OBJ_SENTINEL = u'\u200C'
+CAST_SENTINEL = u'\u24EA'
+
+
+def join_key(*args):
+    nargs = []
+    for arg in args:
+        if not isinstance(arg, six.string_types):
+            arg = CAST_SENTINEL + pickle.dumps(arg)
+        nargs.append(arg)
+    return SEP_SENTINEL.join(nargs)
+
+
+def true_key(key):
+    key = key.rsplit(SEP_SENTINEL, 1)[-1]
+    if key.startswith(CAST_SENTINEL):
+        return pickle.loads(key)
+    return key
+
+
+class BaseProvider(object):
+    def __init__(self, config):
+        self.config = config
+        self.data = {}
+
+    def exists(self, key):
+        return key in self.data
+
+    def keys(self, other):
+        count = other.count(SEP_SENTINEL) + 1
+        for key in self.data.keys():
+            if key.startswith(other) and key.count(SEP_SENTINEL) == count:
+                yield key
+
+    def get_many(self, keys):
+        for key in keys:
+            yield key, self.get(key)
+
+    def get(self, key):
+        return self.data[key]
+
+    def set(self, key, value):
+        self.data[key] = value
+
+    def delete(self, key):
+        del self.data[key]
+
+    def load(self):
+        pass
+
+    def save(self):
+        pass
+
+    def root(self):
+        return StorageDict(self)
+
+
+class StorageDict(UserDict):
+    def __init__(self, parent_or_provider, key=None):
+        if isinstance(parent_or_provider, BaseProvider):
+            self.provider = parent_or_provider
+            self.parent = None
+        else:
+            self.parent = parent_or_provider
+            self.provider = self.parent.provider
+        self._key = key or ROOT_SENTINEL
+
+    def keys(self):
+        return map(true_key, self.provider.keys(self.key))
+
+    def values(self):
+        for key in self.keys():
+            yield self.provider.get(key)
+
+    def items(self):
+        for key in self.keys():
+            yield (true_key(key), self.provider.get(key))
+
+    def ensure(self, key, typ=dict):
+        if key not in self:
+            self[key] = typ()
+        return self[key]
+
+    def update(self, obj):
+        for k, v in six.iteritems(obj):
+            self[k] = v
+
+    @property
+    def data(self):
+        obj = {}
+
+        for raw, value in self.provider.get_many(self.provider.keys(self.key)):
+            key = true_key(raw)
+
+            if value == OBJ_SENTINEL:
+                value = self.__class__(self, key=key).data
+            obj[key] = value
+        return obj
+
+    @property
+    def key(self):
+        if self.parent is not None:
+            return join_key(self.parent.key, self._key)
+        return self._key
+
+    def __setitem__(self, key, value):
+        if isinstance(value, dict):
+            obj = self.__class__(self, key)
+            obj.update(value)
+            value = OBJ_SENTINEL
+
+        self.provider.set(join_key(self.key, key), value)
+
+    def __getitem__(self, key):
+        res = self.provider.get(join_key(self.key, key))
+
+        if res == OBJ_SENTINEL:
+            return self.__class__(self, key)
+
+        return res
+
+    def __delitem__(self, key):
+        return self.provider.delete(join_key(self.key, key))
+
+    def __contains__(self, key):
+        return self.provider.exists(join_key(self.key, key))

disco/bot/providers/disk.py

@@ -0,0 +1,53 @@
+import os
+import gevent
+
+from disco.util.serializer import Serializer
+from .base import BaseProvider
+
+
+class DiskProvider(BaseProvider):
+    def __init__(self, config):
+        super(DiskProvider, self).__init__(config)
+        self.format = config.get('format', 'pickle')
+        self.path = config.get('path', 'storage') + '.' + self.format
+        self.fsync = config.get('fsync', False)
+        self.fsync_changes = config.get('fsync_changes', 1)
+
+        self.change_count = 0
+
+    def autosave_loop(self, interval):
+        while True:
+            gevent.sleep(interval)
+            self.save()
+
+    def _on_change(self):
+        if self.fsync:
+            self.change_count += 1
+
+            if self.change_count >= self.fsync_changes:
+                self.save()
+                self.change_count = 0
+
+    def load(self):
+        if not os.path.exists(self.path):
+            return
+
+        if self.config.get('autosave', True):
+            self.autosave_task = gevent.spawn(
+                self.autosave_loop,
+                self.config.get('autosave_interval', 120))
+
+        with open(self.path, 'r') as f:
+            self.data = Serializer.loads(self.format, f.read())
+
+    def save(self):
+        with open(self.path, 'w') as f:
+            f.write(Serializer.dumps(self.format, self.data))
+
+    def set(self, key, value):
+        super(DiskProvider, self).set(key, value)
+        self._on_change()
+
+    def delete(self, key):
+        super(DiskProvider, self).delete(key)
+        self._on_change()

disco/bot/providers/memory.py

@@ -0,0 +1,5 @@
+from .base import BaseProvider
+
+
+class MemoryProvider(BaseProvider):
+    pass