Search code examples
pythonargparse

Pythonic way for adding parameters to argparse based on another parameter


In the main function, I have a parser which validates optional inputs:

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--platform',required=True)
    parser.add_argument('--foo')
    parser.add_argument('--bar')
    parser.add_argument('--baz')
    parser.parse_args()

The above snippet is an example which only works when I supply --platform and either --foo,--bar or --baz.

This code is used by different components, let's call them components A, B and C.

Component A actually only specifies --foo and --bar:

python script.py --platform A --foo first_example --bar first_example

Component B actually only specifies --bar and --baz:

python script.py --platform B --bar second_example --baz second_exmaple

Component C actually only specifies --baz:

python script.py --platform C --baz third_example

As I introduce more components, which provide different arguments, the number of arguments I have to add to the parser increases. The above is just an example and I am currently dealing with 20 or so arguments (will likely be more in the future).

I have been thinking about having a configuration file (.yaml) where I define which arguments each component needs:

# parameters.yaml
A:
  - foo
  - bar

B:
  - bar
  - baz

C:
  - baz

I would like to simplify the main function to look at the --platform argument and, based on which platform has been passed as argument, read the configuration and add additional arguments to the parser.

Here's what I have tried:

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--platform',required=True)
    
    # Read from .yaml file
    with open('parameters.yaml') as parameter_file:
        parameters = yaml.safe_load(parameter_file)
    
    for argument in parameters[sys.argv[sys.argv.index('--platform') + 1]]:
        parser.add_argument(
           '--' + argument
    )

    parser.parse_args()

Calling the function:

python script.py --platform C --baz third_example

the above code works but I am looking for other Pythonic solutions as I am a beginner with Python. I don't really like having to look at sys.argv to determine what --platform has been specified. Are there any better solutions to this problem?


Solution

  • With subparsers as sub-commands:

    import argparse
    
    def run_command(parser, args):
        if args.command == 'A':
            print(args)
        elif args.command == 'B':
            print(args)
        elif args.command == 'C':
            print(args)
    
    parser = argparse.ArgumentParser(
        prog='PROG', 
        epilog="See '<command> --help' to read about a specific sub-command."
    )
    subparsers = parser.add_subparsers(dest='command', help='Sub-commands')
    
    A_parser = subparsers.add_parser('A', help='Platform A')
    A_parser.add_argument("--foo")
    A_parser.add_argument('--bar')
    A_parser.set_defaults(func=run_command)
    
    B_parser = subparsers.add_parser('B', help='Platform B')
    B_parser.add_argument('--bar')
    B_parser.add_argument('--baz')
    B_parser.set_defaults(func=run_command)
    
    C_parser = subparsers.add_parser('C', help='Platform C')
    C_parser.add_argument('--baz')
    C_parser.set_defaults(func=run_command)
    
    args = parser.parse_args()
    if args.command is not None:
        args.func(parser, args)
    else:
        parser.print_help()
    

    This generates:

    ~ python args.py -h
    usage: PROG [-h] {A,B,C} ...
    
    positional arguments:
      {A,B,C}     Sub-commands
        A         Platform A
        B         Platform B
        C         Platform C
    
    optional arguments:
      -h, --help  show this help message and exit
    
    See '<command> --help' to read about a specific sub-command.
    

    and

    ~ python args.py B -h
    usage: PROG B [-h] [--bar BAR] [--baz BAZ]
    
    optional arguments:
      -h, --help  show this help message and exit
      --bar BAR
      --baz BAZ