#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File : argument.py
# Author : Jiayuan Mao
# Email : maojiayuan@gmail.com
# Date : 01/18/2018
#
# This file is part of Jacinle.
# Distributed under terms of the MIT license.
"""A customized argument parser."""
import os.path as osp
import argparse
import json
from jacinle.logging import get_logger
from jacinle.utils.enum import JacEnum
from jacinle.utils.printing import print_to
from jacinle.utils.container import G
from .device import DeviceNameFormat, parse_and_set_devices
from .keyboard import str2bool, str2bool_long
__all__ = ['JacArgumentParser']
logger = get_logger(__file__)
[docs]
class JacArgumentParser(argparse.ArgumentParser):
"""A customized argument parser. The main difference is that :class:`JacArgumentParser` supports
additional types, including:
- ``type='bool'``: a boolean value, internally converted to ``True`` (true, t, yes, y, 1) or ``False`` (no, n, false, f, 0).
- ``type='checked_file'``: a file path, checked to be an existing file.
- ``type='checked_dir'``: a directory path, checked to be an existing directory.
- ``type='ensured_dir'``: a directory path, checked to be an existing directory, and created if not.
- ``type='kv'``: a key-value pair, separated by ``=`` and ``;``. For example, ``--configs "data.int_or_float=int_value; data.string='string_value'"``.
- ``action='set_device'``: set CUDA visible devices.
- ``action='as_enum'``: convert the argument to an enum value.
"""
[docs]
def __init__(self, *args, **kwargs):
kwargs.setdefault('fromfile_prefix_chars', '@')
kwargs.setdefault('formatter_class', argparse.ArgumentDefaultsHelpFormatter)
super().__init__(*args, **kwargs)
self.register('type', 'bool', _type_bool)
self.register('type', 'checked_file', _type_checked_file)
self.register('type', 'checked_dir', _type_checked_dir)
self.register('type', 'ensured_dir', _type_ensured_dir)
self.register('type', 'kv', _type_kv)
self.register('action', 'set_device', SetDeviceAction)
self.register('action', 'as_enum', AsEnumAction)
def _type_bool(string):
try:
return str2bool(string)
except ValueError:
raise argparse.ArgumentTypeError()
def _type_checked_file(string):
if not osp.isfile(string):
raise argparse.ArgumentTypeError('Check file existence failed: "{}".'.format(string))
return string
def _type_checked_dir(string):
if not osp.isdir(string):
raise argparse.ArgumentTypeError('Check directory existence failed: "{}".'.format(string))
return string
def _type_ensured_dir(string):
if not osp.isdir(string):
# TODO(Jiayuan Mao @ 05/08): change to a Y/N question.
import jacinle.io as io
io.mkdir(string)
return string
class _KV(object):
def __init__(self, string):
self.string = string
if len(self.string) > 0:
kvs = list(string.split(';'))
else:
kvs = []
for i, kv in enumerate(kvs):
k, v = kv.split('=')
try:
bool_value = str2bool_long(v)
v = bool_value
kvs[i] = (k, v)
continue
except ValueError:
pass
if v.startswith('"') or v.startswith("'"):
assert v.endswith('"') or v.endswith("'")
v = v[1:-1]
try:
v = float(v)
if int(v) == v:
v = int(v)
except:
pass
kvs[i] = (k, v)
self.kvs = kvs
def apply(self, configs):
with print_to(logger.info):
print('Applying KVs:')
for k, v in self.kvs:
print(' kv.{} = {}'.format(k, v))
keys = k.split('.')
current = configs
for k in keys[:-1]:
try:
current = getattr(current, k)
except AttributeError:
current = current.setdefault(k, G())
try:
setattr(current, keys[-1], v)
except AttributeError:
current[keys[-1]] = v
def __jsonify__(self):
return json.dumps(self.kvs)
def _type_kv(string):
"""
In the format of:
--configs "data.int_or_float=int_value; data.string='string_value'"
"""
return _KV(string)
[docs]
class SetDeviceAction(argparse.Action):
[docs]
def __init__(
self, option_strings, dest, format='int', set_device=True, nargs=None, const=None, default=None,
type=None, choices=None, required=False, help=None, metavar=None
):
DeviceNameFormat.assert_valid(format)
self.format = format
self.set_device = set_device
super().__init__(
option_strings=option_strings, dest=dest, nargs=nargs, const=const, default=default,
type=type, choices=choices, required=required, help=help, metavar=metavar
)
[docs]
def __call__(self, parser, namespace, values, option_string=None):
setattr(namespace, self.dest, parse_and_set_devices(values, self.format, self.set_device))
[docs]
class AsEnumAction(argparse.Action):
[docs]
def __init__(
self, option_strings, dest, type, nargs=None, const=None, default=None, choices=None,
required=False, help=None, metavar=None
):
assert issubclass(type, JacEnum)
self.enum_type = type
if choices is None:
choices = type.choice_values()
if default is not None:
default = self.enum_type.from_string(default)
super().__init__(
option_strings=option_strings, dest=dest, nargs=nargs, const=const, default=default,
type=None, choices=choices, required=required, help=help, metavar=metavar
)
[docs]
def __call__(self, parser, namespace, values, option_string=None):
if isinstance(values, (tuple, list)):
setattr(namespace, self.dest, tuple(map(self.enum_type.from_string, values)))
else:
setattr(namespace, self.dest, self.enum_type.from_string(values))