Allow add_multi_constructor with Noneperlpunk/default-multi

Loader.add_multi_constructor(None, myconstructor)

Also add test for add_multi_constructor('!', ...) etc.

See issue #317

8 files changed, 137 insertions, 2 deletions

diff --git a/lib/yaml/constructor.py b/lib/yaml/constructor.py
index 094b3e0..951ae69 100644
--- a/lib/yaml/constructor.py
+++ b/lib/yaml/constructor.py
@@ -74,7 +74,7 @@ class BaseConstructor(object):
             constructor = self.yaml_constructors[node.tag]
         else:
             for tag_prefix in self.yaml_multi_constructors:
-                if node.tag.startswith(tag_prefix):
+                if tag_prefix is not None and node.tag.startswith(tag_prefix):
                     tag_suffix = node.tag[len(tag_prefix):]
                     constructor = self.yaml_multi_constructors[tag_prefix]
                     break
diff --git a/lib3/yaml/constructor.py b/lib3/yaml/constructor.py
index f8d71a1..b96ff5e 100644
--- a/lib3/yaml/constructor.py
+++ b/lib3/yaml/constructor.py
@@ -72,7 +72,7 @@ class BaseConstructor:
             constructor = self.yaml_constructors[node.tag]
         else:
             for tag_prefix in self.yaml_multi_constructors:
-                if node.tag.startswith(tag_prefix):
+                if tag_prefix is not None and node.tag.startswith(tag_prefix):
                     tag_suffix = node.tag[len(tag_prefix):]
                     constructor = self.yaml_multi_constructors[tag_prefix]
                     break
diff --git a/tests/data/multi-constructor.code b/tests/data/multi-constructor.code
new file mode 100644
index 0000000..590d852
--- /dev/null
+++ b/tests/data/multi-constructor.code
@@ -0,0 +1,4 @@
+[
+    {'Tag1': ['a', 1, 'b', 2]},
+    {'Tag2': ['a', 1, 'b', 2]},
+]
diff --git a/tests/data/multi-constructor.multi b/tests/data/multi-constructor.multi
new file mode 100644
index 0000000..4f5d71f
--- /dev/null
+++ b/tests/data/multi-constructor.multi
@@ -0,0 +1,3 @@
+---
+- !Tag1 [a, 1, b, 2]
+- !!Tag2 [a, 1, b, 2]
diff --git a/tests/lib/test_multi_constructor.py b/tests/lib/test_multi_constructor.py
new file mode 100644
index 0000000..f6e28fe
--- /dev/null
+++ b/tests/lib/test_multi_constructor.py
@@ -0,0 +1,63 @@
+import yaml
+import pprint
+import sys
+
+def _load_code(expression):
+    return eval(expression)
+
+def myconstructor1(constructor, tag, node):
+    seq = constructor.construct_sequence(node)
+    return {tag: seq }
+
+def myconstructor2(constructor, tag, node):
+    seq = constructor.construct_sequence(node)
+    string = ''
+    try:
+        i = tag.index('!') + 1
+    except:
+        try:
+            i = tag.rindex(':') + 1
+        except:
+            pass
+    if i >= 0:
+        tag = tag[i:]
+    return { tag: seq }
+
+class Multi1(yaml.FullLoader):
+    pass
+class Multi2(yaml.FullLoader):
+    pass
+
+def test_multi_constructor(input_filename, code_filename, verbose=False):
+    input = open(input_filename, 'rb').read().decode('utf-8')
+    native = _load_code(open(code_filename, 'rb').read())
+
+    # default multi constructor for ! and !! tags
+    Multi1.add_multi_constructor('!', myconstructor1)
+    Multi1.add_multi_constructor('tag:yaml.org,2002:', myconstructor1)
+
+    data = yaml.load(input, Loader=Multi1)
+    if verbose:
+        print('Multi1:')
+        print(data)
+        print(native)
+    assert(data == native)
+
+
+    # default multi constructor for all tags
+    Multi2.add_multi_constructor(None, myconstructor2)
+
+    data = yaml.load(input, Loader=Multi2)
+    if verbose:
+        print('Multi2:')
+        print(data)
+        print(native)
+    assert(data == native)
+
+
+test_multi_constructor.unittest = ['.multi', '.code']
+
+if __name__ == '__main__':
+    import test_appliance
+    test_appliance.run(globals())
+
diff --git a/tests/lib/test_yaml.py b/tests/lib/test_yaml.py
index e26342d..352cd8d 100644
--- a/tests/lib/test_yaml.py
+++ b/tests/lib/test_yaml.py
@@ -12,6 +12,7 @@ from test_representer import *
 from test_recursive import *
 from test_input_output import *
 from test_sort_keys import *
+from test_multi_constructor import *
 
 if __name__ == '__main__':
     import test_appliance
diff --git a/tests/lib3/test_multi_constructor.py b/tests/lib3/test_multi_constructor.py
new file mode 100644
index 0000000..f6e28fe
--- /dev/null
+++ b/tests/lib3/test_multi_constructor.py
@@ -0,0 +1,63 @@
+import yaml
+import pprint
+import sys
+
+def _load_code(expression):
+    return eval(expression)
+
+def myconstructor1(constructor, tag, node):
+    seq = constructor.construct_sequence(node)
+    return {tag: seq }
+
+def myconstructor2(constructor, tag, node):
+    seq = constructor.construct_sequence(node)
+    string = ''
+    try:
+        i = tag.index('!') + 1
+    except:
+        try:
+            i = tag.rindex(':') + 1
+        except:
+            pass
+    if i >= 0:
+        tag = tag[i:]
+    return { tag: seq }
+
+class Multi1(yaml.FullLoader):
+    pass
+class Multi2(yaml.FullLoader):
+    pass
+
+def test_multi_constructor(input_filename, code_filename, verbose=False):
+    input = open(input_filename, 'rb').read().decode('utf-8')
+    native = _load_code(open(code_filename, 'rb').read())
+
+    # default multi constructor for ! and !! tags
+    Multi1.add_multi_constructor('!', myconstructor1)
+    Multi1.add_multi_constructor('tag:yaml.org,2002:', myconstructor1)
+
+    data = yaml.load(input, Loader=Multi1)
+    if verbose:
+        print('Multi1:')
+        print(data)
+        print(native)
+    assert(data == native)
+
+
+    # default multi constructor for all tags
+    Multi2.add_multi_constructor(None, myconstructor2)
+
+    data = yaml.load(input, Loader=Multi2)
+    if verbose:
+        print('Multi2:')
+        print(data)
+        print(native)
+    assert(data == native)
+
+
+test_multi_constructor.unittest = ['.multi', '.code']
+
+if __name__ == '__main__':
+    import test_appliance
+    test_appliance.run(globals())
+
diff --git a/tests/lib3/test_yaml.py b/tests/lib3/test_yaml.py
index e26342d..352cd8d 100644
--- a/tests/lib3/test_yaml.py
+++ b/tests/lib3/test_yaml.py
@@ -12,6 +12,7 @@ from test_representer import *
 from test_recursive import *
 from test_input_output import *
 from test_sort_keys import *
+from test_multi_constructor import *
 
 if __name__ == '__main__':
     import test_appliance