workspace.py 8.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279
  1. # Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
  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. from __future__ import absolute_import
  15. from __future__ import print_function
  16. from __future__ import division
  17. import importlib
  18. import os
  19. import sys
  20. import yaml
  21. import copy
  22. import collections
  23. try:
  24. collectionsAbc = collections.abc
  25. except AttributeError:
  26. collectionsAbc = collections
  27. from .config.schema import SchemaDict, SharedConfig, extract_schema
  28. from .config.yaml_helpers import serializable
  29. __all__ = [
  30. 'global_config',
  31. 'load_config',
  32. 'merge_config',
  33. 'get_registered_modules',
  34. 'create',
  35. 'register',
  36. 'serializable',
  37. 'dump_value',
  38. ]
  39. def dump_value(value):
  40. # XXX this is hackish, but collections.abc is not available in python 2
  41. if hasattr(value, '__dict__') or isinstance(value, (dict, tuple, list)):
  42. value = yaml.dump(value, default_flow_style=True)
  43. value = value.replace('\n', '')
  44. value = value.replace('...', '')
  45. return "'{}'".format(value)
  46. else:
  47. # primitive types
  48. return str(value)
  49. class AttrDict(dict):
  50. """Single level attribute dict, NOT recursive"""
  51. def __init__(self, **kwargs):
  52. super(AttrDict, self).__init__()
  53. super(AttrDict, self).update(kwargs)
  54. def __getattr__(self, key):
  55. if key in self:
  56. return self[key]
  57. raise AttributeError("object has no attribute '{}'".format(key))
  58. global_config = AttrDict()
  59. BASE_KEY = '_BASE_'
  60. # parse and load _BASE_ recursively
  61. def _load_config_with_base(file_path):
  62. with open(file_path) as f:
  63. file_cfg = yaml.load(f, Loader=yaml.Loader)
  64. # NOTE: cfgs outside have higher priority than cfgs in _BASE_
  65. if BASE_KEY in file_cfg:
  66. all_base_cfg = AttrDict()
  67. base_ymls = list(file_cfg[BASE_KEY])
  68. for base_yml in base_ymls:
  69. if base_yml.startswith("~"):
  70. base_yml = os.path.expanduser(base_yml)
  71. if not base_yml.startswith('/'):
  72. base_yml = os.path.join(os.path.dirname(file_path), base_yml)
  73. with open(base_yml) as f:
  74. base_cfg = _load_config_with_base(base_yml)
  75. all_base_cfg = merge_config(base_cfg, all_base_cfg)
  76. del file_cfg[BASE_KEY]
  77. return merge_config(file_cfg, all_base_cfg)
  78. return file_cfg
  79. def load_config(file_path):
  80. """
  81. Load config from file.
  82. Args:
  83. file_path (str): Path of the config file to be loaded.
  84. Returns: global config
  85. """
  86. _, ext = os.path.splitext(file_path)
  87. assert ext in ['.yml', '.yaml'], "only support yaml files for now"
  88. # load config from file and merge into global config
  89. cfg = _load_config_with_base(file_path)
  90. cfg['filename'] = os.path.splitext(os.path.split(file_path)[-1])[0]
  91. merge_config(cfg)
  92. return global_config
  93. def dict_merge(dct, merge_dct):
  94. """ Recursive dict merge. Inspired by :meth:``dict.update()``, instead of
  95. updating only top-level keys, dict_merge recurses down into dicts nested
  96. to an arbitrary depth, updating keys. The ``merge_dct`` is merged into
  97. ``dct``.
  98. Args:
  99. dct: dict onto which the merge is executed
  100. merge_dct: dct merged into dct
  101. Returns: dct
  102. """
  103. for k, v in merge_dct.items():
  104. if (k in dct and isinstance(dct[k], dict) and
  105. isinstance(merge_dct[k], collectionsAbc.Mapping)):
  106. dict_merge(dct[k], merge_dct[k])
  107. else:
  108. dct[k] = merge_dct[k]
  109. return dct
  110. def merge_config(config, another_cfg=None):
  111. """
  112. Merge config into global config or another_cfg.
  113. Args:
  114. config (dict): Config to be merged.
  115. Returns: global config
  116. """
  117. global global_config
  118. dct = another_cfg or global_config
  119. return dict_merge(dct, config)
  120. def get_registered_modules():
  121. return {
  122. k: v
  123. for k, v in global_config.items() if isinstance(v, SchemaDict)
  124. }
  125. def make_partial(cls):
  126. op_module = importlib.import_module(cls.__op__.__module__)
  127. op = getattr(op_module, cls.__op__.__name__)
  128. cls.__category__ = getattr(cls, '__category__', None) or 'op'
  129. def partial_apply(self, *args, **kwargs):
  130. kwargs_ = self.__dict__.copy()
  131. kwargs_.update(kwargs)
  132. return op(*args, **kwargs_)
  133. if getattr(cls, '__append_doc__', True): # XXX should default to True?
  134. if sys.version_info[0] > 2:
  135. cls.__doc__ = "Wrapper for `{}` OP".format(op.__name__)
  136. cls.__init__.__doc__ = op.__doc__
  137. cls.__call__ = partial_apply
  138. cls.__call__.__doc__ = op.__doc__
  139. else:
  140. # XXX work around for python 2
  141. partial_apply.__doc__ = op.__doc__
  142. cls.__call__ = partial_apply
  143. return cls
  144. def register(cls):
  145. """
  146. Register a given module class.
  147. Args:
  148. cls (type): Module class to be registered.
  149. Returns: cls
  150. """
  151. if cls.__name__ in global_config:
  152. raise ValueError("Module class already registered: {}".format(
  153. cls.__name__))
  154. if hasattr(cls, '__op__'):
  155. cls = make_partial(cls)
  156. global_config[cls.__name__] = extract_schema(cls)
  157. return cls
  158. def create(cls_or_name, **kwargs):
  159. """
  160. Create an instance of given module class.
  161. Args:
  162. cls_or_name (type or str): Class of which to create instance.
  163. Returns: instance of type `cls_or_name`
  164. """
  165. assert type(cls_or_name) in [type, str
  166. ], "should be a class or name of a class"
  167. name = type(cls_or_name) == str and cls_or_name or cls_or_name.__name__
  168. assert name in global_config and \
  169. isinstance(global_config[name], SchemaDict), \
  170. "the module {} is not registered".format(name)
  171. config = global_config[name]
  172. cls = getattr(config.pymodule, name)
  173. cls_kwargs = {}
  174. cls_kwargs.update(global_config[name])
  175. # parse `shared` annoation of registered modules
  176. if getattr(config, 'shared', None):
  177. for k in config.shared:
  178. target_key = config[k]
  179. shared_conf = config.schema[k].default
  180. assert isinstance(shared_conf, SharedConfig)
  181. if target_key is not None and not isinstance(target_key,
  182. SharedConfig):
  183. continue # value is given for the module
  184. elif shared_conf.key in global_config:
  185. # `key` is present in config
  186. cls_kwargs[k] = global_config[shared_conf.key]
  187. else:
  188. cls_kwargs[k] = shared_conf.default_value
  189. # parse `inject` annoation of registered modules
  190. if getattr(cls, 'from_config', None):
  191. cls_kwargs.update(cls.from_config(config, **kwargs))
  192. if getattr(config, 'inject', None):
  193. for k in config.inject:
  194. target_key = config[k]
  195. # optional dependency
  196. if target_key is None:
  197. continue
  198. if isinstance(target_key, dict) or hasattr(target_key, '__dict__'):
  199. if 'name' not in target_key.keys():
  200. continue
  201. inject_name = str(target_key['name'])
  202. if inject_name not in global_config:
  203. raise ValueError(
  204. "Missing injection name {} and check it's name in cfg file".
  205. format(k))
  206. target = global_config[inject_name]
  207. for i, v in target_key.items():
  208. if i == 'name':
  209. continue
  210. target[i] = v
  211. if isinstance(target, SchemaDict):
  212. cls_kwargs[k] = create(inject_name)
  213. elif isinstance(target_key, str):
  214. if target_key not in global_config:
  215. raise ValueError("Missing injection config:", target_key)
  216. target = global_config[target_key]
  217. if isinstance(target, SchemaDict):
  218. cls_kwargs[k] = create(target_key)
  219. elif hasattr(target, '__dict__'): # serialized object
  220. cls_kwargs[k] = target
  221. else:
  222. raise ValueError("Unsupported injection type:", target_key)
  223. # prevent modification of global config values of reference types
  224. # (e.g., list, dict) from within the created module instances
  225. #kwargs = copy.deepcopy(kwargs)
  226. return cls(**cls_kwargs)