CODEWITHSUNDEEP
pythonargparsepython-unittest
Tim Wescott
Tim Wescott

Reputation: 508

Argparse and unittest.main()

I just added a graph utility to a unittest -- basically, the fully automatic version of the test just does a numerical compare, but I want a human to be able to ask for plots.

Just using argparse, unittest.main() was choking if I used my new argument. What I'm currently doing is checking for that argument, then deleting it from sys.argv which just seems wrong.

Is there a better way to skin this cat?

if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description='Test correction'
    )
    parser.add_argument(
        '--plot-results',
        help='Plot results of cal test',
        action='store_true'
    )
    args = parser.parse_args()

    if args.plot_results:
        while '--plot-results' in sys.argv:
            sys.argv.remove('--plot-results')

    unittest.main()

Upvotes: 2

Views: 500

Answers (1)

chepner
chepner

Reputation: 530843

Argument.parse_known_args is basically your second option: parse the arguments you define, and get back the ones you don't recognize to pass on to unittest.main.

if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description='Test correction'
    )
    parser.add_argument(
        '--plot-results',
        help='Plot results of cal test',
        action='store_true'
    )
    args, remaining = parser.parse_known_args()

    # Restore the script name which parse_known_args() also stripped.
    remaining.insert(0, sys.argv[0])
    unittest.main(argv=remaining)

Probably goes without saying, but don't add any arguments to parser that conflict with the parsers used by unittest itself. These are (mostly?) documented here.

Upvotes: 6

Related Questions