Skip to content

Commit

Permalink
Add test to show how to subclass FullLoader with new blacklist
Browse files Browse the repository at this point in the history
  • Loading branch information
ret2libc committed Mar 10, 2020
1 parent f5657f3 commit d9bb589
Show file tree
Hide file tree
Showing 2 changed files with 22 additions and 1 deletion.
5 changes: 5 additions & 0 deletions tests/data/myfullloader.subclass_blacklist
@@ -0,0 +1,5 @@
- !!python/object/new:yaml.MappingNode
args:
state:
mymethod: test
wrong_method: test2
18 changes: 17 additions & 1 deletion tests/lib3/test_constructor.py
Expand Up @@ -3,6 +3,7 @@
import pprint

import datetime
import re
import yaml.tokens

def execute(code):
Expand All @@ -14,7 +15,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
Expand Down Expand Up @@ -222,6 +223,9 @@ def tzname(self, dt):
def dst(self, dt):
return datetime.timedelta(0)

class MyFullLoader(yaml.FullLoader):
state_blacklist_regexp = re.compile('(^mymethod$|^wrong_.*$)')

today = datetime.date.today()

def _load_code(expression):
Expand Down Expand Up @@ -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__']
Expand Down

0 comments on commit d9bb589

Please sign in to comment.