Add support for ArgumentParser based parsing

disco/bot/command.py

@@ -1,4 +1,5 @@
 import re
+import argparse
 
 from holster.enum import Enum
 
@@ -7,6 +8,7 @@ from disco.util.functional import cached_property
 
 ARGS_REGEX = '(?: ((?:\n|.)*)$|$)'
 ARGS_UNGROUPED_REGEX = '(?: (?:\n|.)*$|$)'
+SPLIT_SPACES_NO_QUOTE = re.compile(r'["|\']([^"\']+)["|\']|(\S+)')
 
 USER_MENTION_RE = re.compile('<@!?([0-9]+)>')
 ROLE_MENTION_RE = re.compile('<@&([0-9]+)>')
@@ -21,6 +23,11 @@ CommandLevels = Enum(
 )
 
 
+class PluginArgumentParser(argparse.ArgumentParser):
+    def error(self, message):
+        raise CommandError(message)
+
+
 class CommandEvent(object):
     """
     An event which is created when a command is triggered. Contains information
@@ -138,6 +145,7 @@ class Command(object):
         self.oob = False
         self.context = {}
         self.metadata = {}
+        self.parser = None
 
         self.update(*args, **kwargs)
 
@@ -151,7 +159,7 @@ class Command(object):
     def get_docstring(self):
         return (self.func.__doc__ or '').format(**self.context)
 
-    def update(self, args=None, level=None, aliases=None, group=None, is_regex=None, oob=False, context=None, **kwargs):
+    def update(self, args=None, level=None, aliases=None, group=None, is_regex=None, oob=False, context=None, parser=False, **kwargs):
         self.triggers += aliases or []
 
         def resolve_role(ctx, rid):
@@ -175,13 +183,14 @@ class Command(object):
         def resolve_guild(ctx, gid):
             return ctx.msg.client.state.guilds.get(gid)
 
-        self.raw_args = args
-        self.args = ArgumentSet.from_string(args or '', {
-            'user': self.mention_type([resolve_user], USER_MENTION_RE, user=True),
-            'role': self.mention_type([resolve_role], ROLE_MENTION_RE),
-            'channel': self.mention_type([resolve_channel], CHANNEL_MENTION_RE, allow_plain=True),
-            'guild': self.mention_type([resolve_guild]),
-        })
+        if args:
+            self.raw_args = args
+            self.args = ArgumentSet.from_string(args, {
+                'user': self.mention_type([resolve_user], USER_MENTION_RE, user=True),
+                'role': self.mention_type([resolve_role], ROLE_MENTION_RE),
+                'channel': self.mention_type([resolve_channel], CHANNEL_MENTION_RE, allow_plain=True),
+                'guild': self.mention_type([resolve_guild]),
+            })
 
         self.level = level
         self.group = group
@@ -190,6 +199,9 @@ class Command(object):
         self.context = context or {}
         self.metadata = kwargs
 
+        if parser:
+            self.parser = PluginArgumentParser(prog=self.name, add_help=False)
+
     @staticmethod
     def mention_type(getters, reg=None, user=False, allow_plain=False):
         def _f(ctx, raw):
@@ -253,20 +265,27 @@ class Command(object):
         bool
             Whether this command was successful
         """
-        if len(event.args) < self.args.required_length:
-            raise CommandError(u'Command {} requires {} arguments (`{}`) passed {}'.format(
-                event.name,
-                self.args.required_length,
-                self.raw_args,
-                len(event.args)
-            ))
-
-        try:
-            parsed_args = self.args.parse(event.args, ctx=event)
-        except ArgumentError as e:
-            raise CommandError(e.message)
+        parsed_kwargs = {}
+
+        if self.args:
+            if len(event.args) < self.args.required_length:
+                raise CommandError(u'Command {} requires {} arguments (`{}`) passed {}'.format(
+                    event.name,
+                    self.args.required_length,
+                    self.raw_args,
+                    len(event.args)
+                ))
+
+            try:
+                parsed_kwargs = self.args.parse(event.args, ctx=event)
+            except ArgumentError as e:
+                raise CommandError(e.message)
+        elif self.parser:
+            event.parser = self.parser
+            parsed_kwargs['args'] = self.parser.parse_args(
+                [i[0] or i[1] for i in SPLIT_SPACES_NO_QUOTE.findall(' '.join(event.args))])
 
         kwargs = {}
         kwargs.update(self.context)
-        kwargs.update(parsed_args)
+        kwargs.update(parsed_kwargs)
         return self.plugin.dispatch('command', self, event, **kwargs)

disco/bot/plugin.py

@@ -12,11 +12,7 @@ from disco.util.logging import LoggingClass
 from disco.bot.command import Command, CommandError
 
 
-class PluginDeco(object):
-    """
-    A utility mixin which provides various function decorators that a plugin
-    author can use to create bound event/command handlers.
-    """
+class BasePluginDeco(object):
     Prio = Priority
 
     # TODO: dont smash class methods
@@ -70,6 +66,7 @@ class PluginDeco(object):
         """
         Creates a new command attached to the function.
         """
+
         return cls.add_meta_deco({
             'type': 'command',
             'args': args,
@@ -123,6 +120,25 @@ class PluginDeco(object):
             'kwargs': kwargs,
         })
 
+    @classmethod
+    def add_argument(cls, *args, **kwargs):
+        """
+        Adds an argument to the argument parser.
+        """
+        return cls.add_meta_deco({
+            'type': 'parser.add_argument',
+            'args': args,
+            'kwargs': kwargs,
+        })
+
+
+class PluginDeco(BasePluginDeco):
+    """
+    A utility mixin which provides various function decorators that a plugin
+    author can use to create bound event/command handlers.
+    """
+    parser = BasePluginDeco
+
 
 class Plugin(LoggingClass, PluginDeco):
     """
@@ -191,7 +207,7 @@ class Plugin(LoggingClass, PluginDeco):
         self._post = {'command': [], 'listener': []}
 
         for member in self.meta_funcs:
-            for meta in member.meta:
+            for meta in reversed(member.meta):
                 self.bind_meta(member, meta)
 
     def bind_meta(self, member, meta):
@@ -205,6 +221,14 @@ class Plugin(LoggingClass, PluginDeco):
         elif meta['type'].startswith('pre_') or meta['type'].startswith('post_'):
             when, typ = meta['type'].split('_', 1)
             self.register_trigger(typ, when, member)
+        elif meta['type'].startswith('parser.'):
+            for command in self.commands:
+                if command.func == member:
+                    getattr(command.parser, meta['type'].split('.', 1)[-1])(
+                        *meta['args'],
+                        **meta['kwargs'])
+        else:
+            raise Exception('unhandled meta type {}'.format(meta))
 
     def handle_exception(self, greenlet, event):
         pass

examples/basic_plugin.py

@@ -66,3 +66,11 @@ class BasicPlugin(Plugin):
             event.msg.reply('```\n{}\n```'.format(
                 '\n'.join(parts)
             ))
+
+    @Plugin.command('test', parser=True)
+    @Plugin.parser.add_argument('-a', '--asdf', help='wow')
+    @Plugin.parser.add_argument('--help', action='store_true')
+    def on_test(self, event, args):
+        if args.help:
+            return event.msg.reply(event.parser.format_help())
+        event.msg.reply(args.asdf)