Restore xml.etree.ElementTree after patch

Restore ``ElementTree`` attribute of ``xml.etree`` module after patching

Closes: #54
Co-authored-by: Marien Zwart <marienz@google.com>
Signed-off-by: Christian Heimes <christian@python.org>

CHANGES.txt

@@ -8,6 +8,7 @@ defusedxml 0.7.0.rc2
 
 - Re-add and deprecate ``defusedxml.cElementTree``
 - Use GitHub Actions instead of TravisCI
+- Restore ``ElementTree`` attribute of ``xml.etree`` module after patching
 
 defusedxml 0.7.0.rc1
 --------------------

README.md

@@ -722,6 +722,7 @@ See <https://www.python.org/psf/license> for licensing details.
 
   - Re-add and deprecate `defusedxml.cElementTree`
   - Use GitHub Actions instead of TravisCI
+  - Restore `ElementTree` attribute of `xml.etree` module after patching
 
 ## defusedxml 0.7.0.rc1
 

defusedxml/ElementTree.py

@@ -45,12 +45,21 @@ def _get_py3_cls():
     cmod = sys.modules.pop(cmodname, None)
 
     sys.modules[cmodname] = None
-    pure_pymod = importlib.import_module(pymodname)
-    if cmod is not None:
-        sys.modules[cmodname] = cmod
-    else:
-        sys.modules.pop(cmodname)
-    sys.modules[pymodname] = pymod
+    try:
+        pure_pymod = importlib.import_module(pymodname)
+    finally:
+        # restore module
+        sys.modules[pymodname] = pymod
+        if cmod is not None:
+            sys.modules[cmodname] = cmod
+        else:
+            sys.modules.pop(cmodname, None)
+        # restore attribute on original package
+        etree_pkg = sys.modules["xml.etree"]
+        if pymod is not None:
+            etree_pkg.ElementTree = pymod
+        elif hasattr(etree_pkg, "ElementTree"):
+            del etree_pkg.ElementTree
 
     _XMLParser = pure_pymod.XMLParser
     _iterparse = pure_pymod.iterparse

tests.py

@@ -6,6 +6,7 @@
 import unittest
 import warnings
 
+from xml.etree import ElementTree as orig_elementtree
 from xml.sax.saxutils import XMLGenerator
 from xml.sax import SAXParseException
 from pyexpat import ExpatError
@@ -208,6 +209,11 @@ def test_aliases(self):
         assert self.module.XMLParser is parser
         assert self.module.XMLParse is parser
 
+    def test_import_order(self):
+        from xml.etree import ElementTree as second_elementtree
+
+        self.assertIs(orig_elementtree, second_elementtree)
+
 
 class TestDefusedcElementTree(TestDefusedElementTree):
     module = cElementTree