Fixed implicit resolver leak yaml#67

bridadan · Aug 17, 2017 · efddcac · efddcac
1 parent 7e026bf
commit efddcac
Show file tree

Hide file tree

Showing 4 changed files with 46 additions and 4 deletions.
diff --git a/lib/yaml/resolver.py b/lib/yaml/resolver.py
@@ -4,6 +4,7 @@
 from error import *
 from nodes import *
 
+import itertools
 import re
 
 class ResolverError(YAMLError):
@@ -146,8 +147,8 @@ def resolve(self, kind, value, implicit):
                 resolvers = self.yaml_implicit_resolvers.get(u'', [])
             else:
                 resolvers = self.yaml_implicit_resolvers.get(value[0], [])
-            resolvers += self.yaml_implicit_resolvers.get(None, [])
-            for tag, regexp in resolvers:
+            none_resolvers = self.yaml_implicit_resolvers.get(None, [])
+            for tag, regexp in itertools.chain(resolvers, none_resolvers):
                 if regexp.match(value):
                     return tag
             implicit = implicit[1]

diff --git a/lib3/yaml/resolver.py b/lib3/yaml/resolver.py
@@ -4,6 +4,7 @@
 from .error import *
 from .nodes import *
 
+import itertools
 import re
 
 class ResolverError(YAMLError):
@@ -146,8 +147,8 @@ def resolve(self, kind, value, implicit):
                 resolvers = self.yaml_implicit_resolvers.get('', [])
             else:
                 resolvers = self.yaml_implicit_resolvers.get(value[0], [])
-            resolvers += self.yaml_implicit_resolvers.get(None, [])
-            for tag, regexp in resolvers:
+            none_resolvers = self.yaml_implicit_resolvers.get(None, [])
+            for tag, regexp in itertools.chain(resolvers, none_resolvers):
                 if regexp.match(value):
                     return tag
             implicit = implicit[1]

diff --git a/tests/lib/test_resolver.py b/tests/lib/test_resolver.py
@@ -1,6 +1,7 @@
 
 import yaml
 import pprint
+import re
 
 def test_implicit_resolver(data_filename, detect_filename, verbose=False):
     correct_tag = None
@@ -21,6 +22,25 @@ def test_implicit_resolver(data_filename, detect_filename, verbose=False):
 
 test_implicit_resolver.unittest = ['.data', '.detect']
 
+def test_implicit_resolver_leak(data_filename, detect_filename, verbose=False):
+    # Add any custom resolver
+    tag, regexp = '!any_resolver', re.compile('AnyResolver')
+    yaml.add_implicit_resolver(tag, regexp)
+
+    none_resolvers = yaml.Dumper.yaml_implicit_resolvers.get(None, [])
+    assert (tag, regexp) in none_resolvers
+
+    old_f_resolvers = yaml.Dumper.yaml_implicit_resolvers.get('f', [])
+    assert (tag, regexp) not in old_f_resolvers
+
+    # Dump at least one ScalarNode
+    yaml.dump(False)
+
+    new_f_resolvers = yaml.Dumper.yaml_implicit_resolvers.get('f', [])
+    assert (tag, regexp) not in new_f_resolvers
+
+test_implicit_resolver_leak.unittest = ['.data', '.path']
+
 def _make_path_loader_and_dumper():
     global MyLoader, MyDumper
 

diff --git a/tests/lib3/test_resolver.py b/tests/lib3/test_resolver.py
@@ -1,6 +1,7 @@
 
 import yaml
 import pprint
+import re
 
 def test_implicit_resolver(data_filename, detect_filename, verbose=False):
     correct_tag = None
@@ -21,6 +22,25 @@ def test_implicit_resolver(data_filename, detect_filename, verbose=False):
 
 test_implicit_resolver.unittest = ['.data', '.detect']
 
+def test_implicit_resolver_leak(data_filename, detect_filename, verbose=False):
+    # Add any custom resolver
+    tag, regexp = '!any_resolver', re.compile('AnyResolver')
+    yaml.add_implicit_resolver(tag, regexp)
+
+    none_resolvers = yaml.Dumper.yaml_implicit_resolvers.get(None, [])
+    assert (tag, regexp) in none_resolvers
+
+    old_f_resolvers = yaml.Dumper.yaml_implicit_resolvers.get('f', [])
+    assert (tag, regexp) not in old_f_resolvers
+
+    # Dump at least one ScalarNode
+    yaml.dump(False)
+
+    new_f_resolvers = yaml.Dumper.yaml_implicit_resolvers.get('f', [])
+    assert (tag, regexp) not in new_f_resolvers
+
+test_implicit_resolver_leak.unittest = ['.data', '.path']
+
 def _make_path_loader_and_dumper():
     global MyLoader, MyDumper