author Dmitriy Morozov <dmitriy@mrzv.org>
Mon, 04 Apr 2011 12:56:15 -0700
changeset 189 a46bf1c7f346
parent 182 f8a82de0540d
parent 187 5e96ab1087da
permissions -rw-r--r--
Merged upstream

# (c) Alexander Solovyov, 2009-2011, under terms of the new BSD License
'''Command line arguments parser

import sys, traceback, getopt, types, textwrap, inspect, os, copy
from itertools import imap

__all__ = ['command', 'dispatch']
__version__ = '2.2'
__author__ = 'Alexander Solovyov'
__email__ = 'piranha@piranha.org.ua'

    import locale
    ENCODING = locale.getpreferredencoding()
    if not ENCODING or ENCODING == 'mac-roman' or 'ascii' in ENCODING.lower():
        ENCODING = 'UTF-8'
except locale.Error:
    ENCODING = 'UTF-8'

def write(text, out=None):
    '''Write output to a given stream (stdout by default)'''
    out = out or sys.stdout
    if isinstance(text, unicode):
        return out.write(text.encode(ENCODING))

def err(text):
    '''Write output to stderr'''
    write(text, out=sys.stderr)


# --------
# Public interface
# --------

def command(options=None, usage=None, name=None, shortlist=False, hide=False):
    '''Decorator to mark function to be used for command line processing.

    All arguments are optional:

     - ``options``: options in format described later. If not supplied,
       will be determined from function.
     - ``usage``: usage string for function, replaces ``%name`` with name
       of program or subcommand. In case if it's subcommand and ``%name``
       is not present, usage is prepended by ``name``
     - ``name``: used for multiple subcommands. Defaults to wrapped
       function name
     - ``shortlist``: if command should be included in shortlist. Used
       only with multiple subcommands
     - ``hide``: if command should be hidden from help listing. Used only
       with multiple subcommands, overrides ``shortlist``

    Options should be a list of 4-tuples in format::

      (shortname, longname, default, help)


     - ``shortname`` is a single letter which can be used then as an option
       specifier on command line (like ``-a``). Will be not used if contains
       falsy value (empty string, for example)
     - ``longname`` - main identificator of an option, can be used as on a
       command line with double dashes (like ``--longname``)
     - ``default`` value for an option, type of it determines how option will be
     - ``help`` string displayed as a help for an option when asked to
    def wrapper(func):
            options_ = list(guess_options(func))
        except TypeError:
            options_ = []
            options_ = options_ + list(options)
        except TypeError:

        name_ = name or func.__name__.replace('_', '-')
        if usage is None:
            usage_ = guess_usage(func, options_)
            usage_ = usage
        prefix = hide and '~' or (shortlist and '^' or '')
        CMDTABLE[prefix + name_] = (func, options_, usage_)

        def help_func(name=None):
            return help_cmd(func, replace_name(usage_, sysname()), options_)
        func.help = help_func

        def inner(*args, **opts):
            # look if we need to add 'help' option
                (True for option in reversed(options_)
                 if option[1] == 'help').next()
            except StopIteration:
                options_.append(('h', 'help', False, 'show help'))

            argv = opts.pop('argv', sys.argv[1:])
            if opts.pop('help', False):
                return func.help()

            if args or opts:
                # no catcher here because this is call from Python
                return call_cmd_regular(func, options_)(*args, **opts)

                opts, args = catcher(lambda: parse(argv, options_), func.help)
            except Abort:
                return -1

            if opts.pop('help', False):
                return func.help()

                return catcher(lambda: call_cmd(name_, func)(*args, **opts),
            except Abort:
                return -1

        return inner
    return wrapper

def dispatch(args=None, cmdtable=None, globaloptions=None,
             middleware=lambda x: x):
    '''Dispatch command arguments based on subcommands.

    - ``args``: list of arguments, default: ``sys.argv[1:]``
    - ``cmdtable``: dict of commands in format described below.
      If not supplied, will use functions decorated with ``@command``.
    - ``globaloptions``: list of options which are applied to all
      commands, will contain ``--help`` option at least.
    - ``middleware``: global decorator for all commands.

    cmdtable format description::

      {'name': (function, options, usage)}

    - ``name`` is the name used on command-line. Can contain aliases
      (separate them with ``|``), pointer to a fact that this command
      should be displayed in short help (start name with ``^``), or to
      a fact that this command should be hidden (start name with ``~``)
    - ``function`` is the actual callable
    - ``options`` is options list in format described in docs
    - ``usage`` is the short string of usage
    args = args or sys.argv[1:]
    cmdtable = cmdtable or CMDTABLE

    globaloptions = globaloptions or []
    globaloptions.append(('h', 'help', False, 'display help'))

    cmdtable['help'] = (help_(cmdtable, globaloptions), [], '[TOPIC]')
    help_func = cmdtable['help'][0]

    autocomplete(cmdtable, args, middleware)

        name, func, args, kwargs = catcher(
            lambda: _dispatch(args, cmdtable, globaloptions),
    except Abort:
        return -1

    if name == '_completion':       # skip middleware
        worker = lambda: call_cmd(name, func)(*args, **kwargs)
        mwfunc = middleware(func)
        depth = func == mwfunc and 1 or 2
        worker = lambda: call_cmd(name, mwfunc, depth=depth)(*args, **kwargs)

        return catcher(worker, help_func)
    except Abort:
        return -1

# --------
# Help
# --------

def help_(cmdtable, globalopts):
    '''Help generator for a command table
    def help_inner(name=None):
        '''Show help for a given help topic or a help overview

        With no arguments, print a list of commands with short help messages.

        Given a command name, print help for that command.
        def helplist():
            hlp = {}
            # determine if any command is marked for shortlist
            shortlist = (name == 'shortlist' and
                         any(imap(lambda x: x.startswith('^'), cmdtable)))

            for cmd, info in cmdtable.items():
                if cmd.startswith('~'):
                    continue  # do not display hidden commands
                if shortlist and not cmd.startswith('^'):
                    continue  # short help contains only marked commands
                cmd = cmd.lstrip('^~')
                doc = pretty_doc_string(info[0])
                hlp[cmd] = doc.strip().splitlines()[0].rstrip()

            hlplist = sorted(hlp)
            maxlen = max(map(len, hlplist))

            write('usage: %s <command> [options]\n' % sysname())
            for cmd in hlplist:
                doc = hlp[cmd]
                write(' %-*s  %s\n' % (maxlen, cmd.split('|', 1)[0], doc))

        if not cmdtable:
            return err('No commands specified!\n')

        if not name or name == 'shortlist':
            return helplist()

        aliases, (cmd, options, usage) = findcmd(name, cmdtable)
        return help_cmd(cmd,
                        replace_name(usage, sysname() + ' ' + aliases[0]),
                        options + globalopts)
    return help_inner

def help_cmd(func, usage, options):
    '''show help for given command

    - ``func``: function to generate help for (``func.__doc__`` is taken)
    - ``usage``: usage string
    - ``options``: options in usual format

    >>> def test(*args, **opts):
    ...     """that's a test command
    ...        you can do nothing with this command"""
    ...     pass
    >>> opts = [('l', 'listen', 'localhost',
    ...          'ip to listen on'),
    ...         ('p', 'port', 8000,
    ...          'port to listen on'),
    ...         ('d', 'daemonize', False,
    ...          'daemonize process'),
    ...         ('', 'pid-file', '',
    ...          'name of file to write process ID to')]
    >>> help_cmd(test, 'test [-l HOST] [NAME]', opts)
    test [-l HOST] [NAME]
    that's a test command
           you can do nothing with this command
     -l --listen     ip to listen on (default: localhost)
     -p --port       port to listen on (default: 8000)
     -d --daemonize  daemonize process
        --pid-file   name of file to write process ID to
    write(usage + '\n')
    doc = pretty_doc_string(func)
    write('\n' + doc.strip() + '\n\n')
    if options:

def help_options(options):
    '''Generator for help on options
    yield 'options:\n\n'
    output = []
    for o in options:
        short, name, default, desc = o[:4]
        if hasattr(default, '__call__'):
            default = default(None)
        default = default and ' (default: %s)' % default or ''
        output.append(('%2s%s' % (short and '-%s' % short,
                                  name and ' --%s' % name),
                       '%s%s' % (desc, default)))

    opts_len = max([len(first) for first, second in output if second] or [0])
    for first, second in output:
        if second:
            # wrap description at 78 chars
            second = textwrap.wrap(second, width=(78 - opts_len - 3))
            pad = '\n' + ' ' * (opts_len + 3)
            yield ' %-*s  %s\n' % (opts_len, first, pad.join(second))
            yield '%s\n' % first

# --------
# Options parsing
# --------

def parse(args, options):
    >>> opts = [('l', 'listen', 'localhost',
    ...          'ip to listen on'),
    ...         ('p', 'port', 8000,
    ...          'port to listen on'),
    ...         ('d', 'daemonize', False,
    ...          'daemonize process'),
    ...         ('', 'pid-file', '',
    ...          'name of file to write process ID to')]
    >>> print parse(['-l', '', '--pi', 'test', 'all'], opts)
    ({'pid_file': 'test', 'daemonize': False, 'port': 8000, 'listen': ''}, ['all'])

    argmap, defmap, state = {}, {}, {}
    shortlist, namelist, funlist = '', [], []

    for o in options:
        # might have the fifth completer element
        short, name, default, comment = o[:4]
        if short and len(short) != 1:
            raise OpsterError(
                'Short option should be only a single character: %s' % short)
        if not name:
            raise OpsterError(
                'Long name should be defined for every option')
        # change name to match Python styling
        pyname = name.replace('-', '_')
        argmap['-' + short] = argmap['--' + name] = pyname
        defmap[pyname] = default

        # copy defaults to state
        if isinstance(default, (list, dict)):
            state[pyname] = copy.copy(default)
        elif hasattr(default, '__call__'):
            state[pyname] = None
            state[pyname] = default

        # getopt wants indication that it takes a parameter
        if not (default is None or default is True or default is False):
            if short:
                short += ':'
            if name:
                name += '='
        if short:
            shortlist += short
        if name:

    opts, args = getopt.gnu_getopt(args, shortlist, namelist)

    # transfer result to state
    for opt, val in opts:
        name = argmap[opt]
        t = type(defmap[name])
        if t is types.FunctionType:
            del funlist[funlist.index(name)]
            state[name] = defmap[name](val)
        elif t is list:
        elif t is dict:
                k, v = val.split('=')
            except ValueError:
                raise ParseError(name, "wrong definition: '%s' "
                                 "(should be in format KEY=VALUE)" % val)
            state[name][k] = v
        elif t in (types.NoneType, types.BooleanType):
            state[name] = not defmap[name]
            state[name] = t(val)

    for name in funlist:
        state[name] = defmap[name](None)

    return state, args

# --------
# Subcommand system
# --------

def _dispatch(args, cmdtable, globalopts):
    '''Dispatch arguments list by a command table
    cmd, func, args, options = cmdparse(args, cmdtable, globalopts)

    if options.pop('help', False):
        return 'help', cmdtable['help'][0], [cmd], {}
    elif not cmd:
        return 'help', cmdtable['help'][0], ['shortlist'], {}

    return cmd, func, args, options

def cmdparse(args, cmdtable, globalopts):
    '''Parse arguments list to find a command, options and arguments
    # command is the first non-option
    cmd = None
    for arg in args:
        if not arg.startswith('-'):
            cmd = arg

    if cmd:

        aliases, info = findcmd(cmd, cmdtable)
        cmd = aliases[0]
        possibleopts = list(info[1])
        possibleopts = []


    options, args = parse(args, possibleopts)
    return cmd, cmd and info[0] or None, args, options

def aliases_(cmdtable_key):
    '''Get aliases from a command table key'''
    return cmdtable_key.lstrip("^~").split("|")

def findpossible(cmd, table):
    Return cmd -> (aliases, command table entry)
    for each matching command.
    choice = {}
    for e in table.keys():
        aliases = aliases_(e)
        found = None
        if cmd in aliases:
            found = cmd
            for a in aliases:
                if a.startswith(cmd):
                    found = a
        if found is not None:
            choice[found] = (aliases, table[e])

    return choice

def findcmd(cmd, table):
    """Return (aliases, command table entry) for command string."""
    choice = findpossible(cmd, table)

    if cmd in choice:
        return choice[cmd]

    if len(choice) > 1:
        clist = choice.keys()
        raise AmbiguousCommand(cmd, clist)

    if choice:
        return choice.values()[0]

    raise UnknownCommand(cmd)

# --------
# Helpers
# --------

def guess_options(func):
    '''Get options definitions from function

    They should be declared in a following way:

    def func(longname=(shortname, default, help)):

    See docstring of ``command()`` for description of those variables.
    args, _, _, defaults = inspect.getargspec(func)
    for name, option in zip(args[-len(defaults):], defaults):
            sname, default, hlp = option[:3]
            completer = option[3] if len(option) > 3 else None
            yield (sname, name.replace('_', '-'), default, hlp, completer)
        except TypeError:

def guess_usage(func, options):
    '''Get usage definition for a function
    usage = '%name '
    if options:
        usage += '[OPTIONS] '
    args, varargs = inspect.getargspec(func)[:2]
    argnum = len(args) - len(options)
    if argnum > 0:
        usage += args[0].upper()
        if argnum > 1:
            usage += 'S'
    elif varargs:
        usage += '[%s]' % varargs.upper()
    return usage

def catcher(target, help_func):
    '''Catches all exceptions and prints human-readable information on them
        return target()
    except UnknownCommand, e:
        err("unknown command: '%s'\n" % e)
        raise Abort()
    except AmbiguousCommand, e:
        err("command '%s' is ambiguous:\n    %s\n" %
            (e.args[0], ' '.join(e.args[1])))
        raise Abort()
    except ParseError, e:
        err('%s: %s\n\n' % (e.args[0], e.args[1].strip()))
        raise Abort()
    except getopt.GetoptError, e:
        err('error: %s\n\n' % e)
        raise Abort()
    except OpsterError, e:
        err('%s\n' % e)
        raise Abort()

def call_cmd(name, func, depth=1):
    '''Wrapper for command call, catching situation with insufficient arguments

    ``depth`` is necessary when there is a middleware in setup
    def inner(*args, **kwargs):
            return func(*args, **kwargs)
        except TypeError:
            if len(traceback.extract_tb(sys.exc_info()[2])) == depth:
                raise ParseError(name, "invalid arguments")
    return inner

def call_cmd_regular(func, opts):
    '''Wrapper for command for handling function calls from Python
    def inner(*args, **kwargs):
        funcargs, _, varkw, defaults = inspect.getargspec(func)
        if len(args) > len(funcargs):
            raise TypeError('You have supplied more positional arguments'
                            ' than applicable')

        funckwargs = dict((lname.replace('-', '_'), default)
                          for _, lname, default, _ in opts)
        if 'help' not in (defaults or ()) and not varkw:
            funckwargs.pop('help', None)
        return func(*args, **funckwargs)
    return inner

def replace_name(usage, name):
    '''Replace name placeholder with a command name'''
    if '%name' in usage:
        return usage.replace('%name', name, 1)
    return name + ' ' + usage

def sysname():
    '''Returns name of executing file'''
    name = sys.argv[0]
    if name.startswith('/'):
        return name.rsplit('/', 1)[1]
    elif name.startswith('./'):
        return name[2:]
    return name

def pretty_doc_string(item):
    "Doc string with adjusted indentation level of the 2nd line and beyond."
    raw_doc = item.__doc__ or '(no help text available)'
    lines = raw_doc.strip().splitlines()
    if len(lines) <= 1:
        return raw_doc
    indent = len(lines[1]) - len(lines[1].lstrip())
    return '\n'.join([lines[0]] + map(lambda l: l[indent:], lines[1:]))

    from functools import wraps
except ImportError:
    def wraps(wrapped, assigned=('__module__', '__name__', '__doc__'),
        '''functools.wraps replacement for Python 2.4
        def inner(wrapper):
            for attr in assigned:
                setattr(wrapper, attr, getattr(wrapped, attr))
            for attr in updated:
                getattr(wrapper, attr).update(getattr(wrapped, attr, {}))
            return wrapper
        return inner

# --------
# Autocomplete system
# --------

# Borrowed from PIP
def autocomplete(cmdtable, args, middleware):
    """Command and option completion.

    Enable by sourcing one of the completion shell scripts (bash or zsh).

    # Don't complete if user hasn't sourced bash_completion file.
    if 'OPSTER_AUTO_COMPLETE' not in os.environ:
    cwords = os.environ['COMP_WORDS'].split()[1:]
    cword = int(os.environ['COMP_CWORD'])

        current = cwords[cword - 1]
    except IndexError:
        current = ''

    commands = []
    for k in cmdtable.keys():
        commands += aliases_(k)

    # command
    if cword == 1:
        print ' '.join(filter(lambda x: x.startswith(current), commands))

    # command options
            aliases, (cmd, opts, usage) = findcmd(cwords[0], cmdtable)
        except AmbiguousCommand:

        idx = -2 if current else -1
        options = []

        for o in opts:
            short, long, default, help = o[:4]
            completer = o[4] if len(o) > 4 else None
            short, long = '-%s' % short, '--%s' % long
            options += [short, long]

            if cwords[idx] in (short, long) and completer:
                args = middleware(completer)(current)
                print ' '.join(args),

        print ' '.join((o for o in options if o.startswith(current)))


# opster bash completion start
                   COMP_CWORD=$COMP_CWORD \\
                   OPSTER_AUTO_COMPLETE=1 $1 ) )
complete -o default -F _opster_completion %s
# opster bash completion end
# opster zsh completion start
function _opster_completion {
  local words cword
  read -Ac words
  read -cn cword
  reply=( $( COMP_WORDS="$words[*]" \\
             COMP_CWORD=$(( cword-1 )) \\
             OPSTER_AUTO_COMPLETE=1 $words[1] ) )
compctl -K _opster_completion %s
# opster zsh completion end

@command(name='_completion', hide=True)
def completion(type=('t', 'bash', 'Completion type (bash or zsh)'),
               # kwargs will catch every global option, which we get
               # anyway, because middleware is skipped
    """Outputs completion script for bash or zsh."""

    prog_name = os.path.split(sys.argv[0])[1]
    print COMPLETIONS[type].strip() % prog_name

# --------
# Exceptions
# --------

# Command exceptions
class OpsterError(Exception):
    'Base opster exception'

class AmbiguousCommand(OpsterError):
    'Raised if command is ambiguous'

class UnknownCommand(OpsterError):
    'Raised if command is unknown'

class ParseError(OpsterError):
    'Raised on error in command line parsing'

class Abort(OpsterError):
    'Processing error, abort execution'

if __name__ == '__main__':
    import doctest