580 lines
21 KiB
Python
580 lines
21 KiB
Python
#===============================================================================
|
|
# Copyright 2010 Matt Chaput
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
#===============================================================================
|
|
|
|
import re, sys
|
|
from inspect import getargspec
|
|
from textwrap import wrap
|
|
|
|
|
|
def normalize_docstring(docstring):
|
|
"""Normalizes whitespace in the given string.
|
|
"""
|
|
return re.sub(r"[\r\n\t ]+", " ", docstring).strip()
|
|
|
|
|
|
param_exp = re.compile(r"^([\t ]*):param (.*?): ([^\n]*\n(\1[ \t]+[^\n]*\n)*)",
|
|
re.MULTILINE)
|
|
|
|
def find_param_docs(docstring):
|
|
"""Finds ReStructuredText-style ":param:" lines in the docstring and
|
|
returns a dictionary mapping param names to doc strings.
|
|
"""
|
|
|
|
paramdocs = {}
|
|
for match in param_exp.finditer(docstring):
|
|
name = match.group(2)
|
|
value = match.group(3)
|
|
paramdocs[name] = value
|
|
return paramdocs
|
|
|
|
def remove_param_docs(docstring):
|
|
"""Finds ReStructuredText-style ":param:" lines in the docstring and
|
|
returns a new string with the param documentation removed.
|
|
"""
|
|
return param_exp.sub("", docstring)
|
|
|
|
|
|
def process_docstring(docstring):
|
|
"""Takes a docstring and returns a list of strings representing
|
|
the paragraphs in the docstring.
|
|
"""
|
|
|
|
lines = docstring.split("\n")
|
|
paras = [[]]
|
|
for line in lines:
|
|
if not line.strip():
|
|
paras.append([])
|
|
else:
|
|
paras[-1].append(line)
|
|
paras = [normalize_docstring(" ".join(ls))
|
|
for ls in paras if ls]
|
|
return paras
|
|
|
|
|
|
def format_paras(paras, width, indent=0):
|
|
"""Takes a list of paragraph strings and formats them into a word-wrapped,
|
|
optionally indented string.
|
|
"""
|
|
|
|
output = []
|
|
for para in paras:
|
|
lines = wrap(para, width-indent)
|
|
if lines:
|
|
for line in lines:
|
|
output.append((" " * indent) + line)
|
|
output.append("")
|
|
return "\n".join(output)
|
|
|
|
|
|
def totype(v, default):
|
|
"""Tries to convert the value 'v' into the same type as 'default'.
|
|
"""
|
|
# Python3 does not have long, so fake it
|
|
try:
|
|
long
|
|
except NameError:
|
|
long = int
|
|
|
|
t = type(default)
|
|
if t is int:
|
|
return int(v)
|
|
elif t is float:
|
|
return float(v)
|
|
elif t is long:
|
|
return long(v)
|
|
elif t is bool:
|
|
lv = v.lower()
|
|
if lv in ("true", "yes", "on", "1"):
|
|
return True
|
|
elif lv in ("false", "no", "off", "0"):
|
|
return False
|
|
else:
|
|
raise TypeError
|
|
else:
|
|
return v
|
|
|
|
|
|
class CommandError(Exception): pass
|
|
|
|
|
|
class Cmd(object):
|
|
"""Stores metadata about a command.
|
|
"""
|
|
|
|
def __init__(self, name, fn, argnames, keywords, shortopts,
|
|
has_varargs, has_kwargs, docstring, paramdocs):
|
|
self.name = name
|
|
self.fn = fn
|
|
self.argnames = argnames
|
|
self.keywords = keywords
|
|
self.shortopts = shortopts
|
|
self.has_varargs = has_varargs
|
|
self.has_kwargs = has_kwargs
|
|
self.docstring = docstring
|
|
self.paramdocs = paramdocs
|
|
self.param_keywords = False
|
|
|
|
|
|
class Baker(object):
|
|
def __init__(self):
|
|
self.commands = {}
|
|
self.defaultcommand = None
|
|
|
|
def command(self, fn=None, name=None, default=False,
|
|
params=None, shortopts=None):
|
|
"""Registers a command with the bakery. This does not call the
|
|
function, it simply adds it to the list of functions this Baker
|
|
knows about.
|
|
|
|
This method is usually used as a decorator::
|
|
|
|
b = Baker()
|
|
|
|
@b.command
|
|
def test():
|
|
pass
|
|
|
|
:param fn: the function to register.
|
|
:param name: use this argument to register the command under a
|
|
different name than the function name.
|
|
:param default: if True, this command is used when a command is not
|
|
specified on the command line.
|
|
:param params: a dictionary mapping parameter names to docstrings. If
|
|
you don't specify this argument, parameter annotations will be used
|
|
(Python 3.x only), or the functions docstring will be searched for
|
|
Sphinx-style ':param' blocks.
|
|
:param shortopts: a dictionary mapping parameter names to short
|
|
options, e.g. {"verbose": "v"}.
|
|
"""
|
|
|
|
# This method works as a decorator with or without arguments.
|
|
if fn is None:
|
|
# The decorator was given arguments, e.g. @command(default=True),
|
|
# so we have to return a function that will wrap the function when
|
|
# the decorator is applied.
|
|
return lambda fn: self.command(fn, default=default,
|
|
params=params,
|
|
shortopts=shortopts)
|
|
name = name or fn.__name__
|
|
|
|
# Inspect the argument signature of the function
|
|
arglist, vargsname, kwargsname, defaults = getargspec(fn)
|
|
has_varargs = bool(vargsname)
|
|
has_kwargs = bool(kwargsname)
|
|
|
|
# Get the function's docstring
|
|
docstring = fn.__doc__ or ""
|
|
|
|
# If the user didn't specify parameter help in the decorator
|
|
# arguments, try to get it from parameter annotations (Python 3.x)
|
|
# or RST-style :param: lines in the docstring
|
|
if params is None:
|
|
if hasattr(fn, "func_annotations") and fn.func_annotations:
|
|
params = fn.func_annotations
|
|
else:
|
|
params = find_param_docs(docstring)
|
|
docstring = remove_param_docs(docstring)
|
|
|
|
# If the user didn't specify
|
|
shortopts = shortopts or {}
|
|
|
|
if defaults:
|
|
# Zip up the keyword argument names with their defaults
|
|
keywords = dict(list(zip(arglist[0-len(defaults):], defaults)))
|
|
elif has_kwargs:
|
|
# Allow keyword arguments specified by params.
|
|
keywords = dict(list(zip(params.keys(), [""]*len(params))))
|
|
# But set a flag to detect this
|
|
self.param_keywords = True
|
|
else:
|
|
keywords = {}
|
|
|
|
# If this is a method, remove 'self' from the argument list
|
|
if arglist and arglist[0] == "self":
|
|
arglist.pop(0)
|
|
|
|
# Create a Cmd object to represent this command and store it
|
|
cmd = Cmd(name, fn, arglist, keywords, shortopts,
|
|
has_varargs, has_kwargs,
|
|
docstring, params)
|
|
self.commands[cmd.name] = cmd
|
|
|
|
# If default is True, set this as the default command
|
|
if default:
|
|
self.defaultcommand = cmd
|
|
|
|
return fn
|
|
|
|
def print_top_help(self, scriptname, file=sys.stdout):
|
|
"""Prints the documentation for the script and exits.
|
|
|
|
:param scriptname: the name of the script being executed (argv[0]).
|
|
:param file: the file to write the help to. The default is stdout.
|
|
"""
|
|
|
|
# Get a sorted list of all command names
|
|
cmdnames = sorted(self.commands.keys())
|
|
|
|
# Calculate the indent for the doc strings by taking the longest
|
|
# command name and adding 3 (one space before the name and two after)
|
|
rindent = max(len(name) for name in cmdnames) + 3
|
|
|
|
# Print the basic help for running a command
|
|
file.write("\nUsage: %s COMMAND <options>\n\n" % scriptname)
|
|
|
|
print("Available commands:\n\n")
|
|
for cmdname in cmdnames:
|
|
# Get the Cmd object for this command
|
|
cmd = self.commands[cmdname]
|
|
|
|
# Calculate the padding necessary to fill from the end of the
|
|
# command name to the documentation margin
|
|
tab = " " * (rindent - (len(cmdname)+1))
|
|
file.write(" " + cmdname + tab)
|
|
|
|
# Get the paragraphs of the command's docstring
|
|
paras = process_docstring(cmd.docstring)
|
|
if paras:
|
|
# Print the first paragraph
|
|
file.write(format_paras([paras[0]], 76, indent=rindent).lstrip())
|
|
|
|
file.write("\n")
|
|
file.write('Use "%s <command> --help" for individual command help.\n' % scriptname)
|
|
sys.exit(0)
|
|
|
|
def print_command_help(self, scriptname, cmd, file=sys.stdout):
|
|
"""Prints the documentation for a specific command and exits.
|
|
|
|
:param scriptname: the name of the script being executed (argv[0]).
|
|
:param cmd: the Cmd object representing the command.
|
|
:param file: the file to write the help to. The default is stdout.
|
|
"""
|
|
|
|
# Print the usage for the command
|
|
file.write("\nUsage: %s %s" % (scriptname, cmd.name))
|
|
|
|
# Print the required and "optional" arguments (where optional
|
|
# arguments are keyword arguments with default None).
|
|
for name in cmd.argnames:
|
|
if name not in cmd.keywords:
|
|
# This is a positional argument
|
|
file.write(" <%s>" % name)
|
|
else:
|
|
# This is a keyword argument, so skip it unless the default is
|
|
# None, in which case treat it like an optional argument.
|
|
if cmd.keywords[name] is None:
|
|
file.write(" [<%s>]" % name)
|
|
|
|
if cmd.has_varargs:
|
|
# This command accepts a variable number of positional arguments
|
|
file.write(" [...]")
|
|
file.write("\n\n")
|
|
|
|
# Print the documentation for this command
|
|
paras = process_docstring(cmd.docstring)
|
|
if paras:
|
|
# Print the first paragraph with no indent (usually just a summary
|
|
# line)
|
|
file.write(format_paras([paras[0]], 76))
|
|
|
|
# Print subsequent paragraphs indented by 4 spaces
|
|
if len(paras) > 1:
|
|
file.write("\n")
|
|
file.write(format_paras(paras[1:], 76, indent=4))
|
|
file.write("\n")
|
|
|
|
# Get a sorted list of keyword argument names
|
|
keynames = sorted([key for key, value in cmd.keywords.items() if value is not None])
|
|
# Print documentation for keyword options
|
|
if keynames:
|
|
file.write("Options:\n\n")
|
|
|
|
# Make formatted headings, e.g. " -k --keyword ", and put them in
|
|
# a list like [(name, heading), ...]
|
|
heads = []
|
|
for keyname in keynames:
|
|
if cmd.keywords[keyname] is None: continue
|
|
|
|
head = " --" + keyname
|
|
if keyname in cmd.shortopts:
|
|
head = " -" + cmd.shortopts[keyname] + head
|
|
head += " "
|
|
heads.append((keyname, head))
|
|
|
|
if heads:
|
|
# Find the length of the longest formatted heading
|
|
rindent = max(len(head) for keyname, head in heads)
|
|
# Pad the headings so they're all as long as the longest one
|
|
heads = [(keyname, head + (" " * (rindent - len(head))))
|
|
for keyname, head in heads]
|
|
|
|
# Print the option docs
|
|
for keyname, head in heads:
|
|
# Print the heading
|
|
file.write(head)
|
|
|
|
# If this parameter has documentation, print it after the
|
|
# heading
|
|
if keyname in cmd.paramdocs:
|
|
paras = process_docstring(cmd.paramdocs.get(keyname, ""))
|
|
file.write(format_paras(paras, 76, indent=rindent).lstrip())
|
|
else:
|
|
file.write("\n")
|
|
file.write("\n")
|
|
|
|
if any((cmd.keywords.get(a) is None) for a in cmd.argnames):
|
|
file.write("(specifying a single hyphen (-) in the argument list means all\n")
|
|
file.write("subsequent arguments are treated as bare arguments, not options)\n")
|
|
file.write("\n")
|
|
|
|
sys.exit(0)
|
|
|
|
def parse_args(self, scriptname, cmd, argv):
|
|
keywords = cmd.keywords
|
|
shortopts = cmd.shortopts
|
|
|
|
# shortopts maps long option names to characters. To look up short
|
|
# options, we need to create a reverse mapping.
|
|
shortchars = dict((v, k) for k, v in shortopts.items())
|
|
|
|
# The *vargs list and **kwargs dict to build up from the command line
|
|
# arguments
|
|
vargs = []
|
|
kwargs = {}
|
|
|
|
while argv:
|
|
# Take the next argument
|
|
arg = argv.pop(0)
|
|
|
|
if arg == "--":
|
|
# All arguments following a double hyphen are treated as
|
|
# positional arguments
|
|
vargs.extend(argv)
|
|
break
|
|
|
|
elif arg == "-":
|
|
# What to do with a bare -? Right now, it's ignored.
|
|
continue
|
|
|
|
elif arg.startswith("--"):
|
|
# Process long option
|
|
|
|
value = None
|
|
if "=" in arg:
|
|
# The argument was specified like --keyword=value
|
|
name, value = arg[2:].split("=", 1)
|
|
default = keywords.get(name)
|
|
try:
|
|
value = totype(value, default)
|
|
except TypeError:
|
|
pass
|
|
else:
|
|
# The argument was not specified with a value, assuming
|
|
# this is a flag
|
|
name = arg[2:]
|
|
default = keywords.get(name)
|
|
|
|
if type(default) is bool:
|
|
# If this option is a boolean, it doesn't need a value;
|
|
# specifying it on the command line means "do the
|
|
# opposite of the default".
|
|
value = not default
|
|
else:
|
|
# The next item in the argument list is the value, i.e.
|
|
# --keyword value
|
|
if not argv or argv[0].startswith("-") or self.param_keywords:
|
|
# There isn't a value available, or its derived
|
|
# from parameters. Just use True, assuming this
|
|
# is a flag.
|
|
value = True
|
|
else:
|
|
value = argv.pop(0)
|
|
try:
|
|
value = totype(value, default)
|
|
except TypeError:
|
|
pass
|
|
|
|
# Store this option
|
|
kwargs[name] = value
|
|
|
|
elif arg.startswith("-") and cmd.shortopts:
|
|
# Process short option(s)
|
|
|
|
# For each character after the '-'...
|
|
for i in range(1, len(arg)):
|
|
char = arg[i]
|
|
if char not in shortchars:
|
|
continue
|
|
|
|
# Get the long option name corresponding to this char
|
|
name = shortchars[char]
|
|
|
|
default = keywords[name]
|
|
if type(default) is bool:
|
|
# If this option is a boolean, it doesn't need a value;
|
|
# specifying it on the command line means "do the
|
|
# opposite of the default".
|
|
kwargs[name] = not default
|
|
else:
|
|
# This option requires a value...
|
|
if self.param_keywords:
|
|
# It's derived from parameters. Just use True,
|
|
# assuming this is a flag.
|
|
value = True
|
|
elif i < len(arg)-1:
|
|
# There are other characters after this one, so
|
|
# the rest of the characters must represent the
|
|
# value (i.e. old-style UNIX option like -Nname)
|
|
value = totype(arg[i+1:], default)
|
|
elif argv:
|
|
# This is the last character in the list, so the
|
|
# next argument on the command line is the value.
|
|
value = argv.pop(0)
|
|
else:
|
|
# There is no value available. Just use True,
|
|
# assuming this is a flag.
|
|
value = True
|
|
|
|
try:
|
|
kwargs[name] = totype(value, default)
|
|
except ValueError:
|
|
raise CommandError("Couldn't convert %s value %r to type %s" % (name, value, type(default)))
|
|
break
|
|
else:
|
|
# This doesn't start with "-", so just add it to the list of
|
|
# positional arguments.
|
|
vargs.append(arg)
|
|
|
|
return vargs, kwargs
|
|
|
|
def parse(self, argv=None):
|
|
"""Parses the command and parameters to call from the list of command
|
|
line arguments. Returns a tuple of (Cmd object, position arg list,
|
|
keyword arg dict).
|
|
|
|
:param argv: the list of options passed to the command line (sys.argv).
|
|
"""
|
|
|
|
if argv is None: argv = sys.argv
|
|
|
|
scriptname = argv[0]
|
|
|
|
if (len(argv) < 2) or (argv[1] == "-h" or argv[1] == "--help"):
|
|
# Print the documentation for the script
|
|
self.print_top_help(scriptname)
|
|
|
|
if argv[1] == "help":
|
|
if len(argv) > 2 and argv[2] in self.commands:
|
|
cmd = self.commands[argv[2]]
|
|
self.print_command_help(scriptname, cmd)
|
|
self.print_top_help(scriptname)
|
|
|
|
if len(argv) > 1 and argv[1] in self.commands:
|
|
# The first argument on the command line (after the script name
|
|
# is the command to run.
|
|
cmd = self.commands[argv[1]]
|
|
|
|
if len(argv) > 2 and (argv[2] == "-h" or argv[2] == "--help"):
|
|
# Print the help for this command and exit
|
|
self.print_command_help(scriptname, cmd)
|
|
|
|
options = argv[2:]
|
|
else:
|
|
# No known command was specified. If there's a default command,
|
|
# use that.
|
|
cmd = self.defaultcommand
|
|
if cmd is None:
|
|
raise CommandError("unknown command `%s' specified" % argv[1])
|
|
|
|
options = argv[1:]
|
|
|
|
# Parse the rest of the arguments on the command line and use them to
|
|
# call the command function.
|
|
args, kwargs = self.parse_args(scriptname, cmd, options)
|
|
return (cmd, args, kwargs)
|
|
|
|
def apply(self, cmd, args, kwargs):
|
|
"""Calls the command function.
|
|
"""
|
|
|
|
# Create a list of positional arguments: arguments that are either
|
|
# required (not in keywords), or where the default is None (taken to be
|
|
# an optional positional argument). This is different from the Python
|
|
# calling convention, which will fill in keyword arguments with extra
|
|
# positional arguments.
|
|
posargs = [a for a in cmd.argnames if cmd.keywords.get(a) is None]
|
|
|
|
if len(args) > len(posargs) and not cmd.has_varargs:
|
|
raise CommandError("Too many arguments to %s: %s" % (cmd.name, " ".join(args)))
|
|
|
|
if cmd.keywords:
|
|
for k in sorted(kwargs.keys()):
|
|
if k not in cmd.keywords:
|
|
raise CommandError("Unknown option --%s" % k)
|
|
|
|
# Rearrange the arguments into the order Python expects
|
|
newargs = []
|
|
newkwargs = dict(kwargs)
|
|
for name in cmd.argnames:
|
|
if args and cmd.keywords.get(name) is None:
|
|
# This argument is required or optional and we have a bare arg
|
|
# to fill it
|
|
newargs.append(args.pop(0))
|
|
elif name not in cmd.keywords and not args:
|
|
# This argument is required but we don't have a bare arg to
|
|
# fill it
|
|
raise CommandError("Required argument '%s' not given" % name)
|
|
else:
|
|
# This is a keyword argument
|
|
newkwargs[name] = kwargs.get(name, cmd.keywords[name])
|
|
newargs.extend(args)
|
|
|
|
return cmd.fn(*newargs, **newkwargs)
|
|
|
|
def run(self, argv=None):
|
|
"""Takes a list of command line arguments, parses it into a command
|
|
name and options, and calls the function corresponding to the command
|
|
with the given arguments.
|
|
|
|
:param argv: the list of options passed to the command line (sys.argv).
|
|
"""
|
|
|
|
return self.apply(*self.parse(argv))
|
|
|
|
def test(self, argv=None):
|
|
"""Takes a list of command line arguments, parses it into a command
|
|
name and options, and returns what the resulting function call would
|
|
look like. This may be useful for testing how command line arguments
|
|
would be passed to your functions.
|
|
|
|
:param argv: the list of options passed to the command line (sys.argv).
|
|
"""
|
|
|
|
cmd, args, kwargs = self.parse(argv)
|
|
result = "%s(%s" % (cmd.name, ",".join(repr(a) for a in args))
|
|
if kwargs:
|
|
kws = ",".join("%s=%r" % (k, v) for k, v in kwargs.items())
|
|
result += "," + kws
|
|
result += ")"
|
|
return result
|
|
|
|
|
|
_baker = Baker()
|
|
command = _baker.command
|
|
run = _baker.run
|
|
test = _baker.test
|
|
help = _baker.print_top_help
|