The simplest way is to take advantage of the type=
option in add_argument
like @hpaulj did although it can be generalized with a factory function:
def argconv(**convs):
def parse_argument(arg):
if arg in convs:
return convs[arg]
else:
msg = "invalid choice: {!r} (choose from {})"
choices = ", ".join(sorted(repr(choice) for choice in convs.keys()))
raise argparse.ArgumentTypeError(msg.format(arg,choices))
return parse_argument
then in lieu of ...
just use type=argconv(A=CONST_A, B=CONST_B)
:
parser.add_argument("--foo", type=argconv(A=CONST_A, B=CONST_B))
And then everything will work as you want it to in your example.
The following is the first answer I posted, it is still valid but isn't nearly as simple as the above solution.
An alternate method is to make a class that inherits from argparse.ArgumentParser
and override parse_args
to modify the result as it is generated:
import argparse
class MappedParser(argparse.ArgumentParser):
mapping = {} #backup if you don't use def_mapping
def def_mapping(self,**options):
self.mapping = options
def parse_args(self,args=None,namespace=None):
result = argparse.ArgumentParser.parse_args(self,args,namespace)
for name,options in self.mapping.items(): #by default this is is empty so the loop is skipped
if name in result:
key = getattr(result,name)
if key in options:
replace_with = options[key]
setattr(result,name,replace_with)
else:
self.error("option {name!r} got invalid value: {key!r}\n must be one of {valid}".format(name=name,key=key,valid=tuple(options.keys())))
return #error should exit program but I'll leave this just to be safe.
return result
this way the rest of your (example) program would look like this:
# There is nothing restricting their type.
CONST_A = "<something>"
CONST_B = ["other value", "type is irrelevent"]
parser = MappedParser() #constructor is same
parser.def_mapping(foo={"A":CONST_A, "B":CONST_B})
parser.add_argument("--foo") # and this is unchanged
# the following is now true:
print(parser.parse_args("--foo A".split()).foo is CONST_A)
print(parser.parse_args("--foo B".split()).foo is CONST_B)
#note that 'is' operator works so it is even the same reference
#this gives decent error message
parser.parse_args("--foo INVALID".split())
print("when parser.error() is called the program ends so this never is printed")
Add extra options like this:
parser.def_mapping(foo={"A":CONST_A, "B":CONST_B,"C":"third option"})
or extra arguments like this:
parser.def_mapping(foo={"A":CONST_A, "B":CONST_B},
conv={"int":int,"float":float})
as well any added arguments that are not specified in def_mapping
are left alone so it is very easy to implement.