Source

vanity / vanity / commander.py

Full commit
#!/usr/bin/env python
#
# Copyright 2009  John Mulligan <phlogistonjohn@asynchrono.us>
#
# This software may be used and distributed according to the terms of the
# MIT license, incorporated herein by reference. A copy of this license
# should accompany the source code in a file named COPYING.txt.
#
"""parse command lines with subcommands

This module provides various utilities for hooking multiple
subcommands under a single script. The programmer creates
a table of commands and then either parses the command
line and manually dispatches the command, or uses the main(...)
or launch(...) functions that perform automatic dispatch.

Example:
  >>> # cmd1 and cmd2 are functions
  >>> cmdtable = [
  ...   ('name', cmd1, '', [OPTS], 'something helpful'),
  ...   ('other', cmd2, 'alias|alais2', [], 'something helpless'),
  ...   ]
  >>> ret = commander.launch(OPTS, cmdtable, argv)
  >>> cmd, opts, args = commander.parse(OPTS, cmdtable, argv)

The entries in the command table may be tuples, dictionaries,
or commander.Command objects. As an alternate form of the table
a dictionary mapping command names to dictionaries may be used.

>>> table = {
...     'mycommand': {'target':cmd1, 'opts':[], 'help': 'foo bar'},
...     }


The parse function parses a command line and returns a 
command object, opts, args triple. The launch command parses
and then dispatches the command function, it may raise
a parsing exception or a HelpWanted exception. The programmer
should catch the exception and print usage if this exception is
raised. Finally the main function both dispatches and provides
simple help usage and error handling for cli parsing errors.
By default, error messages will be written to the standard
output.
"""

from vanity import cli
import sys
import traceback


class InvalidCommand(cli.CliError):
    """The CLI was given an invalid command"""
    def __init__(self, value):
        cli.CliError.__init__(self, 'invalid command: %s' % value)


class MissingCommand(cli.CliError):
    """The user failed to specify a command"""
    def __init__(self):
        cli.CliError.__init__(self, 'missing required command name')


class AmbiguousCommand(cli.CliError):
    """The given command was ambigouous"""
    def __init__(self, matches):
        msg = 'possible matches: %s' % ' '.join(matches)
        cli.CliError.__init__(self, msg)
        self.matches = matches


class HelpWanted(cli.CliError):
    """The user has requested help"""
    def __init__(self, cmd, opts, args):
        self.cmd = cmd
        self.opts = opts
        self.args = args


class InvalidArguments(cli.CliError):
    """The user failed to give correct number of args"""
    def __init__(self, name):
        cli.CliError.__init__(self,
            '%s: unexpected number of arguments' % name)


def main(globalopts, cmdtable, arguments, generichelp=None):
    """Launch a subcommand based on the given arguments, automatically
    handling help and cli parsing errors. Will return the
    result of a real command or exit when an error or help
    request is encountered.

    * ``globalopts`` - an options table common to all sub commands
    * ``cmdtable`` -  a commands table
    * ``arguments`` - the command line arguments
    * ``generichelp`` - application help callback function
    """
    try:
        return launch(globalopts, cmdtable, arguments)
    except HelpWanted, herr:
        handlehelp(herr, globalopts, cmdtable)
        sys.exit(0)
    except cli.CliError, err:
        sys.stderr.write('error: %s\n' % err)
        sys.exit(2)


def launch(globalopts, cmdtable, arguments, autohelp=True):
    """Launch a subcommand based on the given arguments, return the
    result of the launched function.

    * ``globalopts`` - an options table common to all sub commands
    * ``cmdtable`` -  a commands table
    * ``arguments`` - the command line arguments
    """
    # inject help objects (if requested)
    cmdtable = CommandTable(cmdtable)
    globalopts = cli.OptionTable(globalopts)
    if autohelp:
        cmdtable._table[HELP.name] = HELP
        globalopts._table[HELPOPT.name] = HELPOPT
    # parse cli
    try:
        cmd, opts, args = parse(globalopts, cmdtable, arguments)
    except MissingCommand:
        cmd = HELP
        opts, args = cli.parse(globalopts, arguments, strict=True)
        if not opts.get('help'):
            raise
    # if help requested: raise HelpWanted
    if cmd == HELP or opts.get('help'):
        raise HelpWanted(cmd, opts, args)
    # fire off the command
    return cmd(opts, args)


def parse(globalopts, cmdtable, arguments):
    """Parse a cli and return the corresponding command, parsed options
    and arguments.

    * ``globalopts`` - an options table common to all sub commands
    * ``cmdtable`` -  a commands table
    * ``arguments`` - the command line arguments
    """
    gopts = EarlyOptionTable(globalopts)
    cmds = CommandTable(cmdtable)
    opts, args = cli._parse(gopts, arguments, True)
    # TODO : handle missing command
    if not args:
        raise MissingCommand()
    cname, args = args[0], args[1:]
    cmd = cmds.find(cname)
    copts, cargs = cli._parse(cmd.opts + gopts, args, cmd.strict)
    copts.update(opts)
    return cmd, copts, cargs


def handlehelp(helperr, globalopts, cmdtable, output=None, generic=None):
    """Handle a HelpWanted excpetion by writing usage information
    into the ``output`` object, which is stdout by default.
    Pass a ``generic`` callback function to produce a generic
    application wide help text. Generic accepts a cmdtable argument
    and must return a generator.

    * ``helperr`` - an instance of HelpWanted
    * ``globalopts`` - an options table common to all sub commands
    * ``cmdtable`` -  a commands table
    * ``output`` - an file-like object opened for writing
    * ``generic`` - application help callback function
    """
    if output is None:
        output = sys.stdout
    if not generic:
        generic = generic_usage
    # determine what type of help we're getting
    if helperr.cmd != HELP:
        check = helperr.cmd.name
    elif helperr.args:
        check = helperr.args[0]
    else:
        check = None
    # produce a help text iter
    cmdtable = CommandTable(cmdtable)
    try:
        if check is None:
            content = generic(cmdtable)
        else:
            content = usage(cmdtable.find(check))
    except cli.CliError, err:
        content = ['error: %s' % err]
    for line in content:
        output.write('%s\n' % line)


def generic_usage(cmdtable):
    """Generic help text generator for all subcommands in the
    application (cmdtable).
    """
    yield 'Application Subcommands:'
    yield ''
    for cmd in sorted(cmdtable):
        for line in usage(cmd, short=True):
            yield '  %-12s  %s' % (cmd.name, line)



def usage(cmd, short=False, prefix=''):
    """Generate command usage.

    * ``cmd`` - generate usage for command
    * ``short`` - generate short usage if true
    * ``prefix`` - string prepended before subcommand name
    """
    if short:
        yield cmd.docstring().splitlines()[0]
        return
    yield 'usage: %s%s [OPTIONS] %s' % (prefix, cmd.name, cmd.help)
    doc = cmd.docstring().strip()
    if doc:
        yield ''
        yield '%s' % doc
        yield ''
    if cmd.opts and list(cmd.opts.longopts()):
        yield 'OPTIONS:'
        for line in cli.usage(cmd.opts):
            yield line


def cleanopts(opts, dropkeys=None):
    """Convert raw options dictionary to one suitable for use a keyword
    arguments to a python function.

    * ``opts`` - The options dictionary to clean
    * ``dropkeys`` - A list of keys to drop if they are in the opts
    """
    opts = dict([(k.replace('-', '_'), v) for k, v in opts.iteritems()])
    if dropkeys:
        for key in dropkeys:
            if key in opts:
                del opts[key]
    return opts


def firstframe():
    """returns true if the type error was caused by the first
    function call in the stack
    """
    return len(traceback.extract_tb(sys.exc_info()[2])) == 1


class Command(object):
    """A thin class representing a command.

    Attributes: name, target, aliases, help, strict
    """

    def __init__(self, name, target, aliases=None, opts=None, 
                 help=None, strict=None):
        self.name = name
        self.target = target
        self.aliases = aliases
        self.help = help
        self.opts = cli.OptionTable(opts)
        self.strict = strict

    def __call__(self, opts, args):
        opts = cleanopts(opts, ('help',))
        try:
            return self.target(*args, **opts)
        except TypeError, err:
            if firstframe():
                raise InvalidArguments(self.name)
            raise

    def aliaslist(self):
        """Return a list of command aliases.
        """
        if self.aliases:
            return self.aliases.split('|')
        else:
            return []

    def docstring(self):
        """Return a documentation string for the command.
        """
        doc = getattr(self.target, '__doc__', None)
        if doc:
            return doc
        else:
            return 'No usage available'
    
    def __cmp__(self, other):
        return cmp(self.name, other.name)

    @classmethod
    def convert(cls, obj):
        """Convert a tuple, dictionary or existing command object
        to a command.
        """
        if isinstance(obj, cls):
            return obj
        if hasattr(obj, 'keys'):
            return cls(**obj)
        return cls(*obj)


class EarlyOptionTable(cli.OptionTable):
    """Special case options table for before-command opts"""

    def assemble(self, opts):
        """an assemble function that will not return all keys"""
        assembled = cli.OptionTable.assemble(self, opts)
        return dict((name, assembled[name]) for name in self._names(opts))

    def _names(self, opts):
        for key, _ in opts:
            try:
                yield self.getlong(key).name
            except KeyError:
                yield self.getshort(key).name


class CommandTable(object):
    """A table of launchable sub-commands.
    """

    def __init__(self, table):
        self._table = {}
        if hasattr(table, 'keys') and hasattr(table, 'items'):
            # if a user is giving us a dict, the entries **must** be dicts
            for name, entry in table.items():
                self._table[name] = Command(name, **entry)
        else:
            for entry in table:
                cmd = Command.convert(entry)
                self._table[cmd.name] = cmd
        return

    def __iter__(self):
        return self._table.itervalues()

    def find(self, name):
        """Find the command entry that best matches the given name.
        Returns a command, if no matches are possible an InvalidCommand
        exception is raised, if multiple matches are possible an
        AmbiguousCommand exception is raised.
        """
        namemap = dict((k,k) for k in self._table)
        for key, cmd in self._table.iteritems():
            for alias in cmd.aliaslist():
                namemap[alias] = key
        if name in namemap:
            key = namemap[name]
            return self._table[key]
        # search partials
        canidates = set()
        for alias in namemap:
            if alias.startswith(name):
                canidates.add(alias)
        if len(canidates) == 1:
            key = namemap[canidates.pop()]
            print namemap
            return self._table[key]
        if not canidates:
            raise InvalidCommand(name)
        else:
            raise AmbiguousCommand(canidates)

    def add(self, name=None, aliases=None, opts=None, help=None):
        """Returns a decorator to add a function to the current table.

        * ``name`` - command name; automatically determined if not given
        * ``aliases`` - a pipe delimited strint of alternate names
        * ``opts`` - an options table; see vanity.cli for details
        * ``help`` - a short command help string
        """
        if not opts:
            opts = []
        if not help:
            help = '[OPTIONS] [ARGS]'
        def _wrap(func):
            if name:
                _name = name
            else:
                _name = func.__name__
            cmd = Command(_name, func, aliases, opts, help)
            self._table[cmd.name] = cmd
            return func
        return _wrap


HELP = Command('help', None, None, [], 'display command help')
HELPOPT = cli.Option('help', 'h', None, 'display command help')