From 2356d355808d2c870cd75b09fa399377f3bf09ff Mon Sep 17 00:00:00 2001
From: xiongqingliang <xiongqingliang@corp.netease.com>
Date: Tue, 8 Aug 2017 19:27:08 +0800
Subject: [PATCH] Fixed implicit resolver leak #67

---
 lib/yaml/resolver.py        |  5 +++--
 lib3/yaml/resolver.py       |  5 +++--
 tests/lib/test_resolver.py  | 20 ++++++++++++++++++++
 tests/lib3/test_resolver.py | 20 ++++++++++++++++++++
 4 files changed, 46 insertions(+), 4 deletions(-)

diff --git a/lib/yaml/resolver.py b/lib/yaml/resolver.py
index 528fbc0e..f721560f 100644
--- 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
index 02b82e73..5310636e 100644
--- 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
index 5566750e..bdc18db5 100644
--- 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
index f059dabb..2e285a1e 100644
--- 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