summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
authorChristian Heimes <christian@python.org>2021-01-12 16:37:17 +0100
committerChristian Heimes <christian@python.org>2021-01-12 18:00:06 +0100
commit3a48453780793c3e98251b0139e9e89e310fb617 (patch)
tree7e2331b416e0fd962a9f04031f7e655f512957e5
parent3010d3f8c81f8e7cbb0d6e102801b09d5d6300d4 (diff)
downloaddefusedxml-git-3a48453780793c3e98251b0139e9e89e310fb617.tar.gz
Restore xml.etree.ElementTree after patch
Restore ``ElementTree`` attribute of ``xml.etree`` module after patching Closes: https://github.com/tiran/defusedxml/issues/54 Co-authored-by: Marien Zwart <marienz@google.com> Signed-off-by: Christian Heimes <christian@python.org>
-rw-r--r--CHANGES.txt1
-rw-r--r--README.md1
-rw-r--r--defusedxml/ElementTree.py21
-rw-r--r--tests.py6
4 files changed, 23 insertions, 6 deletions
diff --git a/CHANGES.txt b/CHANGES.txt
index 3f011de..4671370 100644
--- a/CHANGES.txt
+++ b/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
--------------------
diff --git a/README.md b/README.md
index 33e6b13..a462bf9 100644
--- a/README.md
+++ b/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
diff --git a/defusedxml/ElementTree.py b/defusedxml/ElementTree.py
index b1504e4..55c123e 100644
--- a/defusedxml/ElementTree.py
+++ b/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
diff --git a/tests.py b/tests.py
index f8b9274..89fb985 100644
--- a/tests.py
+++ b/tests.py
@@ -6,6 +6,7 @@ import sys
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 @@ class TestDefusedElementTree(BaseTests):
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