From 6014d391033bd29e50645ade439c58a90c7e8fd2 Mon Sep 17 00:00:00 2001
From: Yuxin Wu <ppwwyyxxc@gmail.com>
Date: Thu, 6 Aug 2020 15:24:07 -0700
Subject: [PATCH] add cfg.set_new_allowed (#44)

---
 yacs/config.py | 14 ++++++++++++++
 yacs/tests.py  |  4 ++++
 2 files changed, 18 insertions(+)

diff --git a/yacs/config.py b/yacs/config.py
index 17400a1..8068543 100644
--- a/yacs/config.py
+++ b/yacs/config.py
@@ -325,6 +325,20 @@ def raise_key_rename_error(self, full_key):
     def is_new_allowed(self):
         return self.__dict__[CfgNode.NEW_ALLOWED]
 
+    def set_new_allowed(self, is_new_allowed):
+        """
+        Set this config (and recursively its subconfigs) to allow merging
+        new keys from other configs.
+        """
+        self.__dict__[CfgNode.NEW_ALLOWED] = is_new_allowed
+        # Recursively set new_allowed state
+        for v in self.__dict__.values():
+            if isinstance(v, CfgNode):
+                v.set_new_allowed(is_new_allowed)
+        for v in self.values():
+            if isinstance(v, CfgNode):
+                v.set_new_allowed(is_new_allowed)
+
     @classmethod
     def load_cfg(cls, cfg_file_obj_or_str):
         """
diff --git a/yacs/tests.py b/yacs/tests.py
index 23bc871..1971754 100644
--- a/yacs/tests.py
+++ b/yacs/tests.py
@@ -305,6 +305,10 @@ def test_new_allowed_bad(self):
         with self.assertRaises(KeyError):
             cfg.merge_from_file("example/config_new_allowed_bad.yaml")
 
+        cfg.set_new_allowed(True)
+        cfg.merge_from_file("example/config_new_allowed_bad.yaml")
+        assert cfg.KWARGS.Y.f == 4
+
 
 class TestCfgNodeSubclass(unittest.TestCase):
     def test_merge_cfg_from_file(self):