Update yacs.py

Author: QiuJueqin

utils/yacs.py

@@ -4,10 +4,11 @@
 # Author: Qiu Jueqin (qiujueqin@gmail.com)
 
 
+import os
 import os.path as op
-import copy
 import argparse
 import pathlib
+from copy import deepcopy
 from collections import OrderedDict, defaultdict
 from contextlib import contextmanager
 
@@ -55,7 +56,7 @@
 
 class Config(OrderedDict):
 
-    def __init__(self, init=None):
+    def __init__(self, init=None, **kwargs):
         """
         :param init: dict | yaml filepath | argparse.Namespace
         """
@@ -69,7 +70,7 @@ def __init__(self, init=None):
         elif isinstance(init, str):
             self.from_yaml(init)
         elif isinstance(init, argparse.Namespace):
-            self.from_namespace(init)
+            self.from_namespace(init, **kwargs)
         else:
             raise TypeError(
                 f'Config could only be instantiated from a dict, a yaml '
@@ -106,10 +107,17 @@ def unfreeze(self):
     def _set_immutable(self, is_immutable):
         """ Recursively set immutability. """
 
-        self.__dict__['__immutable__'] = is_immutable
-        for v in self.values():
-            if isinstance(v, Config):
-                v._set_immutable(is_immutable)
+        def _recursively_set_immutable(obj):
+            if isinstance(obj, dict):
+                if isinstance(obj, Config):
+                    obj.__dict__['__immutable__'] = is_immutable
+                for v in obj.values():
+                    _recursively_set_immutable(v)
+            elif isinstance(obj, (list, tuple)):
+                for item in obj:
+                    _recursively_set_immutable(item)
+
+        _recursively_set_immutable(self)
 
     # ---------------- Set & Get ----------------
 
@@ -163,7 +171,7 @@ def from_yaml(self, yaml_path):
         super().__init__(Config._from_dict(dic))
         self.freeze()
 
-    def from_namespace(self, namespace):
+    def from_namespace(self, parsed_args, unknown_args=None):
         """
         Instantiation from an argparse.Namespace object.
 
@@ -181,40 +189,77 @@ def from_namespace(self, namespace):
 
         Given the returned argparse.Namespace object 'args', from_namespace()
         will create a Config object as if it was instantiated from a nested
-        dict d = {'foo': {'bar': 42}}
+        dict d = {'foo': {'bar': 42}}.
+
+        Optionally, the extra argument `unknown_args` also accepts unknown
+        arguments by parser.parse_known_args(), but note that the arguments
+        in command line must starts with '--'.
+
+        For example, creating an argparse.ArgumentParser with '--foo'
+        argument:
+
+        >>> parser = argparse.ArgumentParser()
+        >>> parser.add_argument('--foo', type=int, default=0)
+
+        but in command line user also inputs other arguments:
+
+        ```
+        python main.py --foo 42 --bar ['Alice', 'Bob']
+        ```
+
+        Given the returned argparse.Namespace object `parsed` and unknown
+        args list `unknown` by calling
+
+         >>> parsed, unknown = parser.parse_known_args()
+
+        `from_namespace(parsed, unknown_args=unknown)` will create a Config
+        object as if it was instantiated from a dict
+        d = {'foo': 42, 'bar': ['Alice', 'Bob']}.
         """
 
-        if not isinstance(namespace, argparse.Namespace):
+        if not isinstance(parsed_args, argparse.Namespace):
             raise TypeError(
                 f'expected an argparse.Namespace object, but given a '
-                f'{type(namespace)} '
+                f'{type(parsed_args)} '
+            )
+
+        nested_dict = self._separator_dict_to_nested_dict(vars(parsed_args))
+
+        if unknown_args:
+            nested_dict.update(
+                self._separator_dict_to_nested_dict(self._unknown_args_to_dict(unknown_args))
             )
 
-        nested_dict = self._separator_dict_to_nested_dict(vars(namespace))
         super().__init__(Config._from_dict(nested_dict))
         self.freeze()
 
-    def merge(self, other, allow_new_attr=False, keep_existed_attr=True):
+    def merge(self, other,
+              exclusive=True,
+              max_exclusive_depth=float('Inf'),
+              keep_existed_attr=True):
         """
         Recursively merge from other object
 
         :param other: Config object | dict | yaml filepath |
             argparse.Namespace object
-        :param allow_new_attr: whether allow to add new attributes
+        :param exclusive: if set to True, merging with new fields is forbidden
 
         Example:
 
         >>> cfg = Config({'optimizer': 'adam'})
-        >>> cfg.merge({'lr': 0.001}, allow_new_attr=True)
+        >>> cfg.merge({'lr': 0.001}, exclusive=False)
         >>> cfg.print()
 
         optimizer: adam
         lr: 0.001
 
-        >>> cfg.merge({'weight_decay': 1E-7}, allow_new_attr=False)
+        >>> cfg.merge({'weight_decay': 1E-7}, exclusive=True)
 
         AttributeError: attempted to add a new attribute: weight_decay
 
+        :param max_exclusive_depth: max depth to prevent from merging new
+            attributes, only valid when exclusive=True. Set to 0 is equal to
+            exclusive=False
         :param keep_existed_attr: whether keep those attributes that are not
             in 'other'. You may wish to trigger this if requires to completely
             replace a child Config object. See example/examples.py: Example 5
@@ -225,15 +270,15 @@ def merge(self, other, allow_new_attr=False, keep_existed_attr=True):
         >>> cfg1 = Config({'foo': {'Alice': 0, 'Bob': 1}})
         >>> cfg2 = cfg1.copy()
         >>> another = {'foo': {'Carol': 42}}
-        >>> cfg1.merge(another, allow_new_attr=True)
+        >>> cfg1.merge(another, exclusive=False)
         >>> cfg1.print()
 
         foo:
             Alice: 0
             Bob: 1
             Carol: 42
 
-        >>> cfg2.merge(another, allow_new_attr=True, keep_existed_attr=False)
+        >>> cfg2.merge(another, exclusive=False, keep_existed_attr=False)
         >>> cfg2.print()
 
         foo:
@@ -249,25 +294,24 @@ def merge(self, other, allow_new_attr=False, keep_existed_attr=True):
                 f'attempted to merge from an unsupported {type(other)} object'
             )
 
-        def _merge(source_cfg, other_cfg, add_new, keep_existed):
+        def _merge(source_cfg, other_cfg, excl, keep_existed, _cur_depth=1):
             """ Recursively merge the new Config object into the source one """
-
-            with source_cfg.unfreeze():
+            with source_cfg.unfreeze(), other_cfg.unfreeze():
                 for k, v in other_cfg.items():
-                    if k not in source_cfg and not add_new:
+                    if k not in source_cfg and excl and _cur_depth <= max_exclusive_depth:
                         raise AttributeError(
-                            f'attempted to add an attribute {k} but it is not '
-                            f'found in the source Config. Set `allow_new_attr` '
-                            f'to True if requires to add new attributes'
+                            f'attempted to merge an attribute `{k}` that is not '
+                            f'found in the source Config. Set `exclusive` to False '
+                            f'if requires to add new attributes'
                         )
 
                     if isinstance(v, Config):
-                        if isinstance(source_cfg.get(k, None), Config):
-                            _merge(source_cfg[k], v, add_new, keep_existed)
+                        if isinstance(source_cfg.get(k), Config):
+                            _merge(source_cfg[k], v, excl, keep_existed, _cur_depth=_cur_depth + 1)
                         else:
                             source_cfg[k] = v
                     else:
-                        source_cfg[k] = copy.deepcopy(v)
+                        source_cfg[k] = deepcopy(v)
 
                 if not keep_existed:
                     source_keys = list(source_cfg.keys())
@@ -276,7 +320,7 @@ def _merge(source_cfg, other_cfg, add_new, keep_existed):
                                 not isinstance(source_cfg[k], Config):
                             source_cfg.remove(k)
 
-        _merge(self, other, allow_new_attr, keep_existed_attr)
+        _merge(self, other, exclusive, keep_existed_attr)
 
     # ---------------- Output ----------------
 
@@ -286,16 +330,20 @@ def to_dict(self, alphabetical=False):
         An inverse method to self.from_dict()
         """
 
-        def _to_dict(config):
-            if alphabetical:
-                config = sorted(config.items(), key=lambda x: x[0])
-            dic = dict(config)
-            for k, v in dic.items():
-                if isinstance(v, Config):
-                    dic[k] = _to_dict(v)
-            return dic
+        def _recursively_to_dict(obj):
+            if isinstance(obj, Config):
+                if alphabetical:
+                    obj = sorted(obj.items(), key=lambda x: x[0])
+                dic = dict(obj)
+                for k, v in dic.items():
+                    dic[k] = _recursively_to_dict(v)
+                return dic
+            elif isinstance(obj, (list, tuple)):
+                return tuple(_recursively_to_dict(item) for item in obj)
+            else:
+                return obj
 
-        return _to_dict(self)
+        return _recursively_to_dict(self)
 
     def to_parser(self):
         """
@@ -329,38 +377,68 @@ def to_parser(self):
 
         return parser
 
-    def dump(self, save_path):
+    def dump(self, save_path, ignored_keys=()):
         """ Dump a Config object into a yaml file """
+        if not save_path.endswith('.yaml'):
+            raise TypeError('only yaml file is supported by dump() method')
+
+        def _serialize(obj):
+            serializable_types = (bool, str, int, float, list, tuple, dict, set, type(None))
+            if not isinstance(obj, serializable_types):
+                return '{} <class \'{}\'>'.format(str(obj), obj.__class__.__name__)
+            elif isinstance(obj, dict):
+                return {k: _serialize(v) for k, v in obj.items()}
+            elif isinstance(obj, (list, tuple)):
+                return [_serialize(item) for item in obj]
+            else:
+                return obj
+
+        serializable_dic = _serialize({
+            k: v for k, v in self.to_dict(alphabetical=True).items() if k not in ignored_keys
+        })
+
+        os.makedirs(op.dirname(save_path), exist_ok=True)
         with open(save_path, 'w') as fp:
-            yaml.safe_dump(self.to_dict(alphabetical=True), fp)
+            yaml.dump(serializable_dic, fp)
 
     def copy(self):
         """ Create a deep copy of the Config object """
-        return Config(copy.deepcopy(self.to_dict()))
+        return Config(deepcopy(self.to_dict()))
 
     # ---------------- Misc ----------------
 
+    def __copy__(self):
+        return self.copy()
+
+    def __deepcopy__(self, memo):
+        return self.copy()
+
     def __repr__(self):
         return self.to_dict().__repr__()
 
     def __str__(self):
         return self.to_dict().__repr__()
 
-    def _format(self):
+    def string(self, alphabetical=False, ignored_keys=(), key_width=30, indent=0):
+        ignored_keys = set(ignored_keys)
 
-        def _to_string(dic, indent=0):
+        def _to_string(dic, idt=indent):
             texts = []
-            for k, v in dic.items():
-                if not isinstance(v, Config):
-                    texts += ['{:<25}{}'.format(' ' * indent + k + ':', str(v))]
+            keys = sorted(dic.keys()) if alphabetical else dic.keys()
+            keys = [k for k in keys if k not in ignored_keys]
+            for k in keys:
+                title = ' ' * idt + str(k) + ':'
+                texts += ['{:<{}}'.format(title, key_width + idt)]
+                if not isinstance(dic[k], Config):
+                    texts[-1] += str(dic[k])
                 else:
-                    texts += [k + ':'] + _to_string(v, indent=indent + 2)
+                    texts += _to_string(dic[k], idt=idt + 2)
             return texts
 
         return '\n'.join(_to_string(self))
 
-    def print(self, streamer=print):
-        streamer(self._format())
+    def print(self, streamer=print, alphabetical=False, ignored_keys=None, key_width=40, indent=0):
+        return streamer(self.string(alphabetical, ignored_keys, key_width, indent))
 
     def remove(self, key):
         """ Remove an attribute by its key. """
@@ -374,15 +452,15 @@ def remove(self, key):
 
         del self[key]
 
-    # ---------------- Private ----------------
+    # ---------------- Helpers ----------------
 
     @classmethod
     def _from_dict(cls, dic):
-        dic = copy.deepcopy(OrderedDict(dic))
+        dic = deepcopy(OrderedDict(dic))
         for k, v in dic.items():
             if isinstance(v, dict):
                 dic[k] = cls(v)
-            elif isinstance(v, list):  # load list as tuple for safety
+            elif isinstance(v, (list, tuple)):  # load list as tuple for safety
                 dic[k] = tuple(cls(x) if isinstance(x, dict) else x for x in v)
 
         return dic
@@ -419,7 +497,6 @@ def _separator_dict_to_nested_dict(separator_dict, separator='.'):
         :param separator_dict: a non-nested dict
         :return: a nested dict
         """
-        separator_dict = copy.deepcopy(separator_dict)
 
         def _init_nested_dict():
             return defaultdict(_init_nested_dict)
@@ -432,7 +509,7 @@ def _default_to_dict(d):
 
         nested_dict = _init_nested_dict()
 
-        for k, v in separator_dict.items():
+        for k, v in deepcopy(separator_dict).items():
             tmp_d = nested_dict
             keys = k.split(separator)
             for sub_key in keys[:-1]:
@@ -473,7 +550,6 @@ def _nested_dict_to_separator_dict(nested_dict, separator='.'):
         :param nested_dict: a regular (optionally nested) dict
         :return: a non-nested dict whose keys contain separators
         """
-        nested_dict = copy.deepcopy(nested_dict)
 
         def _create_separator_dict(x, key='', separator_dict={}):
             if isinstance(x, dict):
@@ -484,4 +560,30 @@ def _create_separator_dict(x, key='', separator_dict={}):
                 separator_dict[key] = x
             return separator_dict
 
-        return _create_separator_dict(nested_dict)
+        return _create_separator_dict(deepcopy(nested_dict))
+
+    @staticmethod
+    def _unknown_args_to_dict(unknown_args):
+        """
+        Convert unknown argument list returned by `parser.parse_known_args()` into a dict
+        :param unknown_args: list of arguments, in which the keys must starts with '--' and
+            the values could be any Python literal expression
+        :return: a non-nested dict
+        """
+        dic = {}
+
+        key, value_lst = None, []
+        for item in unknown_args + ['--']:
+            if item.startswith('--'):
+                if key and value_lst:
+                    literal = ' '.join(value_lst)
+                    try:
+                        dic[key] = eval(literal)
+                    except SyntaxError as e:
+                        raise SyntaxError('invalid argument: --{} {}'.format(key, literal))
+
+                key, value_lst = item.replace('--', ''), []
+            else:
+                value_lst.append(item)
+
+        return dic