Improve argument wrapper

This commit is contained in:
Tim Van Baak 2019-12-31 12:37:40 -08:00
parent 7cf4deac39
commit 7f6035c0c2
2 changed files with 10 additions and 3 deletions

View File

@ -33,7 +33,7 @@ def get_parser(valid_commands):
# whether their argument is an ArgumentParser. # whether their argument is an ArgumentParser.
for name, func in valid_commands.items(): for name, func in valid_commands.items():
# Create the subparser. # Create the subparser.
cmd = subp.add_parser(name) cmd = subp.add_parser(name, description=func.__doc__)
# Delegate subparser setup. # Delegate subparser setup.
func(cmd) func(cmd)
# Store function for later execution # Store function for later execution

View File

@ -1,21 +1,28 @@
# Standard library imports # Standard library imports
from argparse import ArgumentParser as AP from argparse import ArgumentParser as AP
from functools import wraps
def add_argument(*args, **kwargs): def add_argument(*args, **kwargs):
def argument_adder(command): def argument_adder(command):
@wraps(command)
def augmented_command(cmd_args): def augmented_command(cmd_args):
if type(cmd_args) is AP: if type(cmd_args) is AP:
cmd_args.add_argument(*args, **kwargs) cmd_args.add_argument(*args, **kwargs)
else: else:
command(cmd_args) command(cmd_args)
augmented_command.__doc__ = command.__doc__
return augmented_command return augmented_command
return argument_adder return argument_adder
def no_argument(command):
@wraps(command)
def augmented_command(cmd_args):
if type(cmd_args) is not AP:
command(cmd_args)
return augmented_command
@add_argument("--foo", action="store_true") @add_argument("--foo", action="store_true")
def command_a(args): def command_a(args):
"""a docstring""" """a docstring"""
print(args.foo) print(args.foo)