BerfScene / utils /parsing_utils.py
3v324v23's picture
init
2f85de4
raw
history blame
6.36 kB
# python3.7
"""Contains the utility functions for parsing arguments."""
import json
import argparse
import click
__all__ = [
'parse_int', 'parse_float', 'parse_bool', 'parse_index', 'parse_json',
'IntegerParamType', 'FloatParamType', 'BooleanParamType', 'IndexParamType',
'JsonParamType', 'DictAction'
]
def parse_int(arg):
"""Parses an argument to integer.
Support converting string `none` and `null` to `None`.
"""
if arg is None:
return None
if isinstance(arg, str) and arg.lower() in ['none', 'null']:
return None
return int(arg)
def parse_float(arg):
"""Parses an argument to float number.
Support converting string `none` and `null` to `None`.
"""
if arg is None:
return None
if isinstance(arg, str) and arg.lower() in ['none', 'null']:
return None
return float(arg)
def parse_bool(arg):
"""Parses an argument to boolean.
`None` will be converted to `False`.
"""
if isinstance(arg, bool):
return arg
if arg is None:
return False
if arg.lower() in ['1', 'true', 't', 'yes', 'y']:
return True
if arg.lower() in ['0', 'false', 'f', 'no', 'n', 'none', 'null']:
return False
raise ValueError(f'`{arg}` cannot be converted to boolean!')
def parse_index(arg, min_val=None, max_val=None):
"""Parses indices.
If the input is a list or tuple, this function has no effect.
If the input is a string, it can be either a comma separated list of numbers
`1, 3, 5`, or a dash separated range `3 - 10`. Spaces in the string will be
ignored.
Args:
arg: The input argument to parse indices from.
min_val: If not `None`, this function will check that all indices are
equal to or larger than this value. (default: None)
max_val: If not `None`, this function will check that all indices are
equal to or smaller than this field. (default: None)
Returns:
A list of integers.
Raises:
ValueError: If the input is invalid, i.e., neither a list or tuple, nor
a string.
"""
if arg is None or arg == '':
indices = []
elif isinstance(arg, int):
indices = [arg]
elif isinstance(arg, (list, tuple)):
indices = list(arg)
elif isinstance(arg, str):
indices = []
if arg.lower() not in ['none', 'null']:
splits = arg.replace(' ', '').split(',')
for split in splits:
numbers = list(map(int, split.split('-')))
if len(numbers) == 1:
indices.append(numbers[0])
elif len(numbers) == 2:
indices.extend(list(range(numbers[0], numbers[1] + 1)))
else:
raise ValueError(f'Invalid type of input: `{type(arg)}`!')
assert isinstance(indices, list)
indices = sorted(list(set(indices)))
for idx in indices:
assert isinstance(idx, int)
if min_val is not None:
assert idx >= min_val, f'{idx} is smaller than min val `{min_val}`!'
if max_val is not None:
assert idx <= max_val, f'{idx} is larger than max val `{max_val}`!'
return indices
def parse_json(arg):
"""Parses a string-like argument following JSON format.
- `None` arguments will be kept.
- Non-string arguments will be kept.
"""
if not isinstance(arg, str):
return arg
try:
return json.loads(arg)
except json.decoder.JSONDecodeError:
return arg
class IntegerParamType(click.ParamType):
"""Defines a `click.ParamType` to parse integer arguments."""
name = 'int'
def convert(self, value, param, ctx): # pylint: disable=inconsistent-return-statements
try:
return parse_int(value)
except ValueError:
self.fail(f'`{value}` cannot be parsed as an integer!', param, ctx)
class FloatParamType(click.ParamType):
"""Defines a `click.ParamType` to parse float arguments."""
name = 'float'
def convert(self, value, param, ctx): # pylint: disable=inconsistent-return-statements
try:
return parse_float(value)
except ValueError:
self.fail(f'`{value}` cannot be parsed as a float!', param, ctx)
class BooleanParamType(click.ParamType):
"""Defines a `click.ParamType` to parse boolean arguments."""
name = 'bool'
def convert(self, value, param, ctx): # pylint: disable=inconsistent-return-statements
try:
return parse_bool(value)
except ValueError:
self.fail(f'`{value}` cannot be parsed as a boolean!', param, ctx)
class IndexParamType(click.ParamType):
"""Defines a `click.ParamType` to parse indices arguments."""
name = 'index'
def __init__(self, min_val=None, max_val=None):
self.min_val = min_val
self.max_val = max_val
def convert(self, value, param, ctx): # pylint: disable=inconsistent-return-statements
try:
return parse_index(value, self.min_val, self.max_val)
except ValueError:
self.fail(
f'`{value}` cannot be parsed as a list of indices!', param, ctx)
class JsonParamType(click.ParamType):
"""Defines a `click.ParamType` to parse arguments following JSON format."""
name = 'json'
def convert(self, value, param, ctx):
return parse_json(value)
class DictAction(argparse.Action):
"""Argparse action to split each argument into (key, value) pair.
Each argument should be with `key=value` format, where `value` should be a
string with JSON format.
For example, with an argparse:
parser.add_argument('--options', nargs='+', action=DictAction)
, you can use following arguments in the command line:
--options \
a=1 \
b=1.5
c=true \
d=null \
e=[1,2,3,4,5] \
f='{"x":1,"y":2,"z":3}' \
NOTE: No space is allowed in each argument. Also, the dictionary-type
argument should be quoted with single quotation marks `'`.
"""
def __call__(self, parser, namespace, values, option_string=None):
options = {}
for argument in values:
key, val = argument.split('=', maxsplit=1)
options[key] = parse_json(val)
setattr(namespace, self.dest, options)