def test_parse_assign_dataclass(self):
flag_values = flags.FlagValues()
def always_fail(v):
raise ValueError()
result = config_flags.DEFINE_config_dataclass('test_config',
_CONFIG,
flag_values=flag_values,
parse_fn=always_fail)
flag_values(['program'])
flag_values['test_config'].value = parse_config_flag('12')
self.assertEqual(result.value.my_model.foo, 12)
def test_flags(default, *flag_args, parse_fn=None):
flag_values = flags.FlagValues()
# DEFINE_config_dataclass accesses sys.argv to build flag list!
old_args = list(sys.argv)
sys.argv[:] = ['', *['--test_config' + f for f in flag_args]]
try:
result = config_flags.DEFINE_config_dataclass('test_config',
default,
flag_values=flag_values,
parse_fn=parse_fn)
_, *remaining = flag_values(sys.argv)
if remaining:
raise ValueError(f'{remaining}')
# assert not remaining
return result.value
finally:
sys.argv[:] = old_args
def test_flag_overrides(self):
# Set up some flag overrides.
old_argv = list(sys.argv)
sys.argv = shlex.split(
'./program foo.py --test_config.baseline_model.foo=99')
flag_values = flags.FlagValues()
# Define a config dataclass flag.
test_config = config_flags.DEFINE_config_dataclass(
'test_config', _CONFIG, flag_values=flag_values)
# Inject the flag overrides.
flag_values(sys.argv)
sys.argv = old_argv
# Did the value get overridden?
self.assertEqual(test_config.value.baseline_model.foo, 99)
_CONFIG = MyConfig(
my_model=MyModelConfig(foo=3,
bar=['a', 'b'],
baz={'foo.b': 'bar'},
buz={
(0, 0): 'ZeroZero',
(0, 1): 'ZeroOne'
}),
baseline_model=MyModelConfig(
foo=55,
bar=['c', 'd'],
),
)
_TEST_FLAG = config_flags.DEFINE_config_dataclass('test_flag', _CONFIG,
'MyConfig data')
def test_flags(default, *flag_args, parse_fn=None):
flag_values = flags.FlagValues()
# DEFINE_config_dataclass accesses sys.argv to build flag list!
old_args = list(sys.argv)
sys.argv[:] = ['', *['--test_config' + f for f in flag_args]]
try:
result = config_flags.DEFINE_config_dataclass('test_config',
default,
flag_values=flag_values,
parse_fn=parse_fn)
_, *remaining = flag_values(sys.argv)
if remaining:
raise ValueError(f'{remaining}')
_CONFIG = MyConfig(
my_model=MyModelConfig(
foo=3,
bar=['a', 'b'],
baz={'foo': 'bar'},
),
baseline_model=MyModelConfig(
foo=55,
bar=['c', 'd'],
),
)
# Define the flag.
_CONFIG_FLAG = config_flags.DEFINE_config_dataclass('config', _CONFIG)
class TypedConfigFlagsTest(absltest.TestCase):
def test_instance(self):
config = _CONFIG_FLAG.value
self.assertIsInstance(config, MyConfig)
self.assertEqual(config.my_model, _CONFIG.my_model)
self.assertEqual(_CONFIG, config)
def test_flag_overrides(self):
# Set up some flag overrides.
old_argv = list(sys.argv)
sys.argv = shlex.split(
'./program foo.py --test_config.baseline_model.foo=99')
from typing import Any, Mapping, Sequence
from absl import app
from ml_collections import config_flags
@dataclasses.dataclass
class MyConfig:
field1: int
field2: str
nested: Mapping[str, Any]
tuple: Sequence[int]
config = MyConfig(
field1=1,
field2='tom',
nested={'field': 2.23},
tuple=(1, 2, 3),
)
_CONFIG = config_flags.DEFINE_config_dataclass('my_config', config)
def main(_):
print(_CONFIG.value)
if __name__ == '__main__':
app.run(main)