Override cleanup_namespaces to preserve xmlns:xs namespace

Author: Carlton Duffett

src/onelogin/saml2/xml_utils.py

@@ -30,7 +30,28 @@ class OneLogin_Saml2_XML(object):
     dump = staticmethod(etree.dump)
     make_root = staticmethod(etree.Element)
     make_child = staticmethod(etree.SubElement)
-    cleanup_namespaces = staticmethod(etree.cleanup_namespaces)
+
+    @staticmethod
+    def cleanup_namespaces(tree_or_element, top_nsmap=None, keep_ns_prefixes=None):
+        """
+        Keeps the xmlns:xs namespace intact when etree.cleanup_namespaces is invoked.
+        :param tree_or_element: An XML tree or element
+        :type tree_or_element: etree.Element
+        :param top_nsmap: A mapping from namespace prefixes to namespace URIs
+        :type top_nsmap: dict
+        :param keep_ns_prefixes: List of prefixes that should not be removed as part of the cleanup
+        :type keep_ns_prefixes: list
+        :returns: An XML tree or element
+        :rtype: etree.Element
+        """
+        all_prefixes_to_keep = [
+            OneLogin_Saml2_Constants.NS_PREFIX_XS
+        ]
+
+        if keep_ns_prefixes:
+            all_prefixes_to_keep = list(set(all_prefixes_to_keep.extend(keep_ns_prefixes)))
+
+        return etree.cleanup_namespaces(tree_or_element, keep_ns_prefixes=all_prefixes_to_keep)
 
     @staticmethod
     def to_string(xml, **kwargs):