update implemenation, tests

2 files changed, 8 insertions, 11 deletions

diff --git a/src/lxml/apihelpers.pxi b/src/lxml/apihelpers.pxi
index 11f4c697..d2d2023f 100644
--- a/src/lxml/apihelpers.pxi
+++ b/src/lxml/apihelpers.pxi
@@ -389,7 +389,8 @@ cdef int _removeUnusedNamespaceDeclarations(xmlNode* c_element, keep_nsmap) exce
         while c_node and c_ns_list_len:
             if c_node.ns:
                 for i in range(c_ns_list_len):
-                    if c_node.ns is c_ns_list[i].ns:
+                    if (c_node.ns is c_ns_list[i].ns or
+                        (keep_nsmap is not None and c_ns_list[i].ns.href in keep_nsmap.values())):
                         c_ns_list_len -= 1
                         c_ns_list[i] = c_ns_list[c_ns_list_len]
                         #c_ns_list[c_ns_list_len] = _ns_node_ref(NULL, NULL)
diff --git a/src/lxml/tests/test_etree.py b/src/lxml/tests/test_etree.py
index fae06192..d200d21e 100644
--- a/src/lxml/tests/test_etree.py
+++ b/src/lxml/tests/test_etree.py
@@ -2687,20 +2687,16 @@ class ETreeOnlyTestCase(HelperTestCase):
             self.etree.tostring(root))
 
     def test_namespace_cleanup_deep_keep_to_top(self):
-        xml = ('<root xmlns:foo="FOO">' +
-               ''.join('<a xmlns:n{n}="NS{n}">'.format(n=i) for i in range(100)) +
-               '<n64:x xmlns:a="A" a:attr="X"/>' +
-               '</a>'*100 +
-               '<foo>foo:bar</foo>' +
+        xml = ('<root xmlns:n64="NS64" xmlns:foo="FOO">' +
+               '<a><n64:x xmlns:a="A" a:attr="X"/>' +
+               '</a><foo>foo:bar</foo>' +
                '</root>').encode('utf8')
         root = self.etree.fromstring(xml)
         self.assertEqual(xml, self.etree.tostring(root))
-        self.etree.cleanup_namespaces(root,
-                                      top_nsmap={'n64': 'NS64'},
-                                      keep_nsmap={'foo': 'FOO'})
+        self.etree.cleanup_namespaces(root, keep_nsmap={'foo': 'FOO'})
         self.assertEqual(
-            b'<root xmlns:n64="NS64" xmlns:foo="FOO">' + b'<a>'*100 +
-            b'<n64:x xmlns:a="A" a:attr="X"/>' + b'</a>'*100 +
+            b'<root xmlns:n64="NS64" xmlns:foo="FOO">' +
+            b'<a><n64:x xmlns:a="A" a:attr="X"/></a>' +
             b'<foo>foo:bar</foo>' + b'</root>',
             self.etree.tostring(root))