diff options
| author | Christian Heimes <christian@python.org> | 2021-01-12 16:37:17 +0100 |
|---|---|---|
| committer | Christian Heimes <christian@python.org> | 2021-01-12 18:00:06 +0100 |
| commit | 3a48453780793c3e98251b0139e9e89e310fb617 (patch) | |
| tree | 7e2331b416e0fd962a9f04031f7e655f512957e5 /defusedxml | |
| parent | 3010d3f8c81f8e7cbb0d6e102801b09d5d6300d4 (diff) | |
| download | defusedxml-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>
Diffstat (limited to 'defusedxml')
| -rw-r--r-- | defusedxml/ElementTree.py | 21 |
1 files changed, 15 insertions, 6 deletions
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 |