utils.py (4365B)
1 2 from typing import Any, Callable, Dict, List, Union 3 import importlib 4 import logging 5 import sys 6 import types 7 8 logger = logging.getLogger(__name__) 9 10 11 LABELS = ["hero", "villain", "victim", "other"] 12 13 class dotdict(dict): 14 """ Dictionary which can be access through var.key instead of var["key"]. """ 15 def __getattr__(self, name: str): 16 if name not in self: 17 raise AttributeError(f"Config key {name} not found") 18 return dotdict(self[name]) if type(self[name]) is dict else self[name] 19 __setattr__ = dict.__setitem__ 20 __delattr__ = dict.__delitem__ 21 22 23 def eval_arg(config: Dict[str, Any], arg: str) -> None: 24 """ 25 Evaluate arg in the context config, and update it. 26 27 The argument is expected to be of the form: 28 (parent.)*key(=value)? 29 If no value is provided, the key is assumed to be a boolean and True is assigned to it. 30 When passing a string argument through the shell, it must be enclosed in quote (like all python string), which usually need to be escaped. 31 """ 32 key: str 33 value: Any 34 if '=' in arg: 35 key, value = arg.split('=', maxsplit=1) 36 try: 37 value = eval(value, config) 38 except NameError: 39 pass 40 config.pop("__builtins__", None) 41 else: 42 key, value = arg, True 43 path: List[str] = key.split('.') 44 for component in path[:-1]: 45 config = config[component] 46 key = path[-1] 47 if isinstance(config, list): 48 key = int(key) 49 config[key] = value 50 51 52 def import_arg(config: Dict[str, Any], arg: str) -> None: 53 """ 54 Load file arg, and update config with its content. 55 56 The file is loaded in an independent context, all the variable defined in the file (even through import) are added to config, with the exception of builtins and whole modules. 57 """ 58 if arg.endswith(".py"): 59 arg = arg[:-3].replace('/', '.') 60 module: types.ModuleType = importlib.import_module(arg) 61 for key, value in vars(module).items(): 62 if key not in module.__builtins__ and not key.startswith("__") and not isinstance(value, types.ModuleType): # pytype: disable=attribute-error 63 config[key] = value 64 65 66 def parse_args() -> dotdict: 67 """ 68 Parse command line arguments and return config dictionary. 69 70 Two kind of argument are supported: 71 - When the argument starts with -- it is evaluated by the eval_arg function 72 - Otherwise the argument is assumed to be a file which is loaded by the import_arg function 73 """ 74 config: Dict[str, Any] = {} 75 for arg in sys.argv[1:]: 76 if arg.startswith("--"): 77 eval_arg(config, arg[2:]) 78 else: 79 import_arg(config, arg) 80 return dotdict(config) 81 82 83 def display_dict(output: Callable[[str], None], input: Dict[str, Any], depth: int = 0) -> None: 84 """ Display nested dictionaries in input using the provided output function. """ 85 for key, value in input.items(): 86 indent = '\t'*depth 87 output(f"{indent}{key}:") 88 if isinstance(value, dict): 89 output('\n') 90 display_dict(output, value, depth+1) 91 else: 92 output(f" {value}\n") 93 94 95 def print_dict(input: Dict[str, Any]) -> None: 96 """ Print dictionary to standard output. """ 97 display_dict(lambda x: print(x, end=""), input) 98 99 100 def log_dict(logger: logging.Logger, input: Dict[str, Any]) -> None: 101 """ Log dictionary to the provided logger. """ 102 class log: 103 buf: str = "" 104 105 def __call__(self, x: str) -> None: 106 self.buf += x 107 if self.buf.endswith('\n'): 108 logger.info(self.buf[:-1]) 109 self.buf = "" 110 display_dict(log(), input) 111 112 113 def flatten_dict(input: Dict[str, Any]) -> Dict[str, Union[bool, int, float, str]]: 114 """ 115 Replace nested dict by dot-separated keys, and cast keys to simple types. 116 117 repr() is used to cast non-base-type to str. 118 """ 119 def impl(result: Dict[str, Union[bool, int, float, str]], input: Dict[str, Any], prefix: str): 120 for key, value in input.items(): 121 if isinstance(value, dict): 122 impl(result, value, f"{key}.") 123 else: 124 result[f"{prefix}{key}"] = value if type(value) in [bool, int, float, str] else repr(value) 125 126 result: Dict[str, Union[bool, int, float, str]] = {} 127 impl(result, input, "") 128 return result