Jake Levi
Jake Levi

Reputation: 1718

Dynamic argument types in Python + argparse

I want to create a CLI using Python and argparse. The CLI should have options to specify a list of values, and also to dynamically specify the type of the values (str, int, float, etc.) in that list (all arguments in the list have the same type). The values in the list must be converted to the specified type.

I have the following baseline implementation, which does work, but if feels a bit clunky, especially when adding more complex types (or even functions which process the input list of arguments). I was wondering if there is a built-in/smoother/more canonical way to do this?

script.py:

import argparse

arg_type_dict = {t.__name__: t for t in [str, int, float]}

def main(
    sweep_arg_type: str,
    sweep_arg_vals: list,
):
    arg_type = arg_type_dict[sweep_arg_type]
    val_list = [arg_type(val_str) for val_str in sweep_arg_vals]

    print(val_list)

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--sweep_arg_vals", required=True, nargs="+")
    parser.add_argument(
        "--sweep_arg_type",
        required=True,
        choices=sorted(arg_type_dict.keys()),
    )
    args = parser.parse_args()

    main(
        args.sweep_arg_type,
        args.sweep_arg_vals,
    )

Usage examples:

python script.py -h
python script.py --sweep_arg_type int    --sweep_arg_vals 0 1 10 -3
python script.py --sweep_arg_type float  --sweep_arg_vals 0 1 10 -3
python script.py --sweep_arg_type float  --sweep_arg_vals 1.2 3.4
python script.py --sweep_arg_type str    --sweep_arg_vals abc def lmnop

Upvotes: -1

Views: 48

Answers (2)

Hai Vu
Hai Vu

Reputation: 40653

I believe the second option in your own solution is the simplest to implement. Note that the arguments are simply JSON values.

import argparse
import json


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--sweep_arg_vals", required=True, type=json.loads)
    args = parser.parse_args()
    print(args.sweep_arg_vals)


if __name__ == "__main__":
    main()

Sample runs

python3 my.py --sweep_arg_vals '[0, 1, 10, -3]'
[0, 1, 10, -3]

python3 my.py --sweep_arg_vals '[1.1, 1.2]'
[1.1, 1.2]

python3 my.py --sweep_arg_vals '["abc", "def", "lmnop"]'
['abc', 'def', 'lmnop']

Notes

  • I use type=json.loads, which does the conversion
  • No need for nargs="+"

Upvotes: 2

Jake Levi
Jake Levi

Reputation: 1718

Option 1

I wrote a small module called argtypes to perform this functionality:

argtypes.py:

def get_types() -> list["ArgType"]:
    return [IntType(), FloatType(), StrType(), IntList()]

def get_dict():
    return {
        arg_type.get_name(): arg_type
        for arg_type in get_types()
    }

def get_choices():
    return sorted(get_dict().keys())

def convert_arg_list(
    arg_type_str: str,
    arg_vals: list,
):
    type_dict = get_dict()
    arg_type = type_dict[arg_type_str]
    typed_arg_vals = [arg_type.convert_arg(val) for val in arg_vals]
    return typed_arg_vals

class ArgType:
    def convert_arg(self, val_str: str):
        raise NotImplementedError

    @classmethod
    def get_name(cls):
        return cls.__name__.replace("Type", "").lower()

class IntType(ArgType):
    def convert_arg(self, val_str):
        return int(val_str)

class FloatType(ArgType):
    def convert_arg(self, val_str):
        return float(val_str)

class StrType(ArgType):
    def convert_arg(self, val_str):
        return str(val_str)

class IntList(ArgType):
    def convert_arg(self, val_str):
        return [int(i) for i in val_str.split(",")]

...

Refactored script.py:

import argparse
import argtypes

def main(
    sweep_arg_type: str,
    sweep_arg_vals: list,
):
    val_list = argtypes.convert_arg_list(sweep_arg_type, sweep_arg_vals)

    print(val_list)

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--sweep_arg_vals", required=True, nargs="+")
    parser.add_argument(
        "--sweep_arg_type",
        required=True,
        choices=argtypes.get_choices(),
    )
    args = parser.parse_args()

    main(
        args.sweep_arg_type,
        args.sweep_arg_vals,
    )

Additional usage examples:

# All usage examples in the original question still work the same as before
python script.py --sweep_arg_type intlist --sweep_arg_vals 0,1,2 10,100,1000 4 3,-2 -5
# >>> [[0, 1, 2], [10, 100, 1000], [4], [3, -2], [-5]]

Option 2

With all the usual caveats of calling eval with user-input, I have found a very simple solution is simply to use parser.add_argument(..., type=eval) and specify the argument list as a string of Python code:

import argparse

def main(
    sweep_arg_vals: list,
):
    print(sweep_arg_vals)
    print(type(sweep_arg_vals), [type(v) for v in sweep_arg_vals])

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--sweep_arg_vals", required=True, type=eval)
    args = parser.parse_args()

    main(
        args.sweep_arg_vals,
    )

Usage examples:

python script.py -h
python script.py --sweep_arg_vals "[0, 1, 10, -3]"
python script.py --sweep_arg_vals "[1.2, 3.4]"
python script.py --sweep_arg_vals "['abc', 'def', 'lmnop']"

Upvotes: 1

Related Questions