You can not select more than 25 topics Topics must start with a chinese character,a letter or number, can include dashes ('-') and can be up to 35 characters long.

config.py 4.3 kB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125
  1. # Copyright 2021 Huawei Technologies Co., Ltd
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. # ============================================================================
  15. """Parse arguments"""
  16. import os
  17. import ast
  18. import argparse
  19. from pprint import pprint, pformat
  20. import yaml
  21. _config_path = "./mr_config.yaml"
  22. class Config:
  23. """
  24. Configuration namespace. Convert dictionary to members.
  25. """
  26. def __init__(self, cfg_dict):
  27. for k, v in cfg_dict.items():
  28. if isinstance(v, (list, tuple)):
  29. setattr(self, k, [Config(x) if isinstance(x, dict) else x for x in v])
  30. else:
  31. setattr(self, k, Config(v) if isinstance(v, dict) else v)
  32. def __str__(self):
  33. return pformat(self.__dict__)
  34. def __repr__(self):
  35. return self.__str__()
  36. def parse_cli_to_yaml(parser, cfg, helper=None, choices=None, cfg_path="default_config.yaml"):
  37. """
  38. Parse command line arguments to the configuration according to the default yaml.
  39. Args:
  40. parser: Parent parser.
  41. cfg: Base configuration.
  42. helper: Helper description.
  43. cfg_path: Path to the default yaml config.
  44. """
  45. parser = argparse.ArgumentParser(description="[REPLACE THIS at config.py]",
  46. parents=[parser])
  47. helper = {} if helper is None else helper
  48. choices = {} if choices is None else choices
  49. for item in cfg:
  50. if not isinstance(cfg[item], list) and not isinstance(cfg[item], dict):
  51. help_description = helper[item] if item in helper else "Please reference to {}".format(cfg_path)
  52. choice = choices[item] if item in choices else None
  53. if isinstance(cfg[item], bool):
  54. parser.add_argument("--" + item, type=ast.literal_eval, default=cfg[item], choices=choice,
  55. help=help_description)
  56. else:
  57. parser.add_argument("--" + item, type=type(cfg[item]), default=cfg[item], choices=choice,
  58. help=help_description)
  59. args = parser.parse_args()
  60. return args
  61. def parse_yaml(yaml_path):
  62. """
  63. Parse the yaml config file.
  64. Args:
  65. yaml_path: Path to the yaml config.
  66. """
  67. with open(yaml_path, 'r') as fin:
  68. try:
  69. cfgs = yaml.load_all(fin.read(), Loader=yaml.FullLoader)
  70. cfgs = [x for x in cfgs]
  71. if len(cfgs) == 1:
  72. cfg_helper = {}
  73. cfg = cfgs[0]
  74. elif len(cfgs) == 2:
  75. cfg, cfg_helper = cfgs
  76. else:
  77. raise ValueError("At most 2 docs (config and help description for help) are supported in config yaml")
  78. print(cfg_helper)
  79. except:
  80. raise ValueError("Failed to parse yaml")
  81. return cfg, cfg_helper
  82. def merge(args, cfg):
  83. """
  84. Merge the base config from yaml file and command line arguments.
  85. Args:
  86. args: Command line arguments.
  87. cfg: Base configuration.
  88. """
  89. args_var = vars(args)
  90. for item in args_var:
  91. cfg[item] = args_var[item]
  92. return cfg
  93. def get_config():
  94. """
  95. Get Config according to the yaml file and cli arguments.
  96. """
  97. parser = argparse.ArgumentParser(description="default name", add_help=False)
  98. current_dir = os.path.dirname(os.path.abspath(__file__))
  99. parser.add_argument("--config_path", type=str, default=os.path.join(current_dir, "../mr_config.yaml"),
  100. help="Config file path")
  101. path_args, _ = parser.parse_known_args()
  102. default, helper = parse_yaml(path_args.config_path)
  103. pprint(default)
  104. args = parse_cli_to_yaml(parser, default, helper, path_args.config_path)
  105. final_config = merge(args, default)
  106. return Config(final_config)
  107. config = get_config()