[project @ Make elementtree importing in openid.yadis.etxrd more robust]

 * Use the importElementTree function in oidutil
 * Handle exceptions from any kind of elementtree library

1 files changed, 23 insertions, 5 deletions

diff --git a/openid/yadis/etxrd.py b/openid/yadis/etxrd.py
index da63cbb..419a968 100644
--- a/openid/yadis/etxrd.py
+++ b/openid/yadis/etxrd.py
@@ -18,17 +18,35 @@ __all__ = [
     'expandServices',
     ]
 
+import sys
 import random
 
-from elementtree.ElementTree import ElementTree
+from openid.oidutil import importElementTree
+ElementTree = importElementTree()
 
 # Use expat if it's present. Otherwise, use xmllib
 try:
-    from xml.parsers.expat import ExpatError as XMLError
-    from elementtree.ElementTree import XMLTreeBuilder
+    XMLTreeBuilder = ElementTree.XMLTreeBuilder
+
+    # This will raise an ImportError if an XML parser is not present.
+    p = XMLTreeBuilder()
 except ImportError:
     from elementtree.SimpleXMLTreeBuilder import TreeBuilder as XMLTreeBuilder
-    from xmllib import Error as XMLError
+
+# the different elementtree modules don't have a common exception
+# model. We just want to be able to catch the exceptions that signify
+# malformed XML data and wrap them, so that the other library code
+# doesn't have to know which XML library we're using.
+try:
+    # Make the parser raise an exception so we can sniff out the type
+    # of exceptions
+    p.feed('> purposely malformed XML <')
+    p.close()
+except (SystemExit, MemoryError, AssertionError, ImportError):
+    raise
+except:
+    XMLError = sys.exc_info()[0]
+    del p
 
 from openid.yadis import xri
 
@@ -64,7 +82,7 @@ def parseXRDS(text):
         exc.reason = why
         raise exc
     else:
-        tree = ElementTree(element)
+        tree = ElementTree.ElementTree(element)
         if not isXRDS(tree):
             raise XRDSError('Not an XRDS document')