diff --git a/disco/bot/command.py b/disco/bot/command.py index c06e30f..fbf7047 100644 --- a/disco/bot/command.py +++ b/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) diff --git a/disco/bot/plugin.py b/disco/bot/plugin.py index 088f4da..155e581 100644 --- a/disco/bot/plugin.py +++ b/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 diff --git a/examples/basic_plugin.py b/examples/basic_plugin.py index 60c5e8b..c57866c 100644 --- a/examples/basic_plugin.py +++ b/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)