How to specify a range using argparse?

Question

My parser states the following:

parser.add_argument(
    '--degree',
    type=int,
    default=3,
    nargs="*",
    help="Degree of the 'poly' kernel. (default=3)."
)

I would like to specify a range instead of just using nargs.

For example: 0 8 2 (start, stop, step) should specify that I want the values 0 2 4 6 8.

Should I just use np.arange as datatype or maybe just specifying 3 args and creating the array myself. Please LMK if there is a better way.

Answer 1

You can set nargs=3 to require 3 values for the parameter.

parser.add_argument(
    '--degree',
    type=int,
    default=3,
    nargs=3,
)

args = parser.parse_args(['--degree', '1', '2', '3'])
numpy_arr = np.arange(*args.degree)

Alternatively, you can pass comma-separated values and use type to convert the value to a range:

def get_np_arange(value):
   try:
       values = [int(i) for i in value.split(',')]
       assert len(values) in (1, 3)
   except (ValueError, AssertionError):
       raise argparse.ArgumentTypeError(
           'Provide a CSV list of 1 or 3 integers'
       )

   # return our value as is if there is only one
   if len(values) == 1:
       return np.array(values)

   # if there are three - return a range
   return np.arange(*values)


parser.add_argument(
    '--degree',
    type=get_np_arange,
)

args = parser.parse_args()
numpy_arr = args.degree
print(numpy_arr)

Example output:

➜  ~ python test.py --degree 1,2    
usage: prog [-h] --degree DEGREE
prog: error: argument --degree: Provide a CSV list of 1 or 3 integers
➜  ~ python test.py --degree a,2,3
usage: prog [-h] --degree DEGREE
prog: error: argument --degree: Provide a CSV list of 1 or 3 integers
➜  ~ python test.py --degree 20
[20]
➜  ~ python test.py --degree 1,10,4
[1 5 9]