miken32
miken32

Reputation: 42715

Check for certain arguments and assign values to an arbitrary variable

I'm writing a program and need to look for one of two arguments set on the command line, and save a value to a single variable based on which one is set.

If I call the program like this:

myprogram -a --foo 123

I want the variable action set to 'a value'. Call it like this:

myprogram -b --foo 123

And action should be set to 'another value'. Call it with neither:

myprogram -c --foo 123

And it should exit with usage info.

Obviously I can do this with some if statements after the fact:

import argparse
parser = argparse.ArgumentParser()
parser.add_argument('-a', action='count')
parser.add_argument('-b', action='count')
parser.add_argument('--foo')
....
args = parser.parse_args()
if args.a == 0 and args.b == 0:
    parser.print_usage()
    sys.exit(1)
if args.a > 0:
    action = 'a value'
elif args.b > 0:
    action = 'another value'

But I'm wondering if argparse can do it for me with less code. From what I've seen in the documentation it's not possible, but Python is very very new to me. Thanks.

Upvotes: 0

Views: 313

Answers (2)

miken32
miken32

Reputation: 42715

I used hapulj's answer, but still ended up with an if statement to check if neither was set. Then I found the ArgumentParser.add_mutually_exclusive_group() function, and ended up with this, which works perfectly.

import argparse
parser = argparse.ArgumentParser()
actiongroup = parser.add_mutually_exclusive_group(required=True)
actiongroup.add_argument('-a', action='store_const', dest='action', const='a value')
actiongroup.add_argument('-b', action='store_const', dest='action', const='another value')
parser.add_argument('--foo')
....
args = parser.parse_args()

Now, arguments -a and -b can't be omitted, and both can't be specified at the same time.

Upvotes: 0

hpaulj
hpaulj

Reputation: 231385

Look at action='store_const'.

In [240]: parser=argparse.ArgumentParser()
In [241]: parser.add_argument('-a', dest='action',
   action='store_const', const='a value')
In [242]: parser.add_argument('-b', dest='action',
   action='store_const', const='another value')
In [243]: parser.add_argument('--foo')

In [244]: parser.parse_args('-a --foo 123'.split())
Out[244]: Namespace(action='a value', foo='123')

In [245]: parser.parse_args('-b --foo 123'.split())
Out[245]: Namespace(action='another value', foo='123')

In [246]: parser.parse_args('-c --foo 123'.split())
usage: ipython3 [-h] [-a] [-b] [--foo FOO] 
ipython3: error: unrecognized arguments: -c
SystemExit: 2

So args.action will have a value' orb value' depending on the argument. Note both -a and -b store to the same dest.

I left -c undefined, so it uses the normal undefined exit with usage. that could be refined.

Defining a c like this would let you do your own exit:

In [247]: parser.add_argument('-c', dest='action', action='store_const', const='exit')

In [248]: args=parser.parse_args('-c --foo 123'.split())
In [249]: if args.action=='exit':parser.print_usage()
usage: ipython3 [-h] [-a] [-b] [--foo FOO] [-c]

If you used action='store_true' instead of 'count' for your -a and -b, you could simplify the if tree to:

if args.a:
    action = 'a value'
elif args.b:
    action = 'another value'
else:
    parser.print_usage()
    sys.exit(1)

Upvotes: 1

Related Questions