Add test to show how to subclass FullLoader with new blacklist

yaml · Mar 17, 2020 · 9091565 · 9091565
1 parent a946ac6
commit 9091565
Show file tree

Hide file tree

Showing 3 changed files with 39 additions and 2 deletions.
diff --git a/tests/data/myfullloader.subclass_blacklist b/tests/data/myfullloader.subclass_blacklist
@@ -0,0 +1,5 @@
+- !!python/object/new:yaml.MappingNode
+  args:
+  state:
+    mymethod: test
+    wrong_method: test2
diff --git a/tests/lib/test_constructor.py b/tests/lib/test_constructor.py
@@ -17,7 +17,7 @@ def _make_objects():
     global MyLoader, MyDumper, MyTestClass1, MyTestClass2, MyTestClass3, YAMLObject1, YAMLObject2,  \
             AnObject, AnInstance, AState, ACustomState, InitArgs, InitArgsWithState,    \
             NewArgs, NewArgsWithState, Reduce, ReduceWithState, Slots, MyInt, MyList, MyDict,  \
-            FixedOffset, today, execute
+            FixedOffset, today, execute, MyFullLoader
 
     class MyLoader(yaml.Loader):
         pass
@@ -235,6 +235,10 @@ def tzname(self, dt):
         def dst(self, dt):
             return datetime.timedelta(0)
 
+    class MyFullLoader(yaml.FullLoader):
+        def get_state_keys_blacklist(self):
+            return super(MyFullLoader, self).get_state_keys_blacklist() + ['^mymethod$', '^wrong_.*$']
+
     today = datetime.date.today()
 
 def _load_code(expression):
@@ -289,6 +293,18 @@ def test_constructor_types(data_filename, code_filename, verbose=False):
 
 test_constructor_types.unittest = ['.data', '.code']
 
+def test_subclass_blacklist_types(data_filename, verbose=False):
+    _make_objects()
+    try:
+        yaml.load(open(data_filename, 'rb').read(), MyFullLoader)
+    except yaml.YAMLError as exc:
+        if verbose:
+            print("%s:" % exc.__class__.__name__, exc)
+    else:
+        raise AssertionError("expected an exception")
+
+test_subclass_blacklist_types.unittest = ['.subclass_blacklist']
+
 if __name__ == '__main__':
     import sys, test_constructor
     sys.modules['test_constructor'] = sys.modules['__main__']

diff --git a/tests/lib3/test_constructor.py b/tests/lib3/test_constructor.py
@@ -14,7 +14,7 @@ def _make_objects():
     global MyLoader, MyDumper, MyTestClass1, MyTestClass2, MyTestClass3, YAMLObject1, YAMLObject2,  \
             AnObject, AnInstance, AState, ACustomState, InitArgs, InitArgsWithState,    \
             NewArgs, NewArgsWithState, Reduce, ReduceWithState, Slots, MyInt, MyList, MyDict,  \
-            FixedOffset, today, execute
+            FixedOffset, today, execute, MyFullLoader
 
     class MyLoader(yaml.Loader):
         pass
@@ -222,6 +222,10 @@ def tzname(self, dt):
         def dst(self, dt):
             return datetime.timedelta(0)
 
+    class MyFullLoader(yaml.FullLoader):
+        def get_state_keys_blacklist(self):
+            return super().get_state_keys_blacklist() + ['^mymethod$', '^wrong_.*$']
+
     today = datetime.date.today()
 
 def _load_code(expression):
@@ -274,6 +278,18 @@ def test_constructor_types(data_filename, code_filename, verbose=False):
 
 test_constructor_types.unittest = ['.data', '.code']
 
+def test_subclass_blacklist_types(data_filename, verbose=False):
+    _make_objects()
+    try:
+        yaml.load(open(data_filename, 'rb').read(), MyFullLoader)
+    except yaml.YAMLError as exc:
+        if verbose:
+            print("%s:" % exc.__class__.__name__, exc)
+    else:
+        raise AssertionError("expected an exception")
+
+test_subclass_blacklist_types.unittest = ['.subclass_blacklist']
+
 if __name__ == '__main__':
     import sys, test_constructor
     sys.modules['test_constructor'] = sys.modules['__main__']