Fix most of the lxml treebuilder errors. This code is still kindof frgile especilly in the fragment case

jgraham · jgraham · commit 9f8929b3c1aa · 2008-03-24T01:11:49.000Z
--HG--
extra : convert_revision : svn%3Aacbfec75-9323-0410-a652-858a13e371e0/trunk%401134
diff --git a/src/html5lib/treewalkers/lxmletree.py b/src/html5lib/treewalkers/lxmletree.py
@@ -7,21 +7,112 @@
 
 from html5lib.constants import voidElements
 
+class Root(object):
+    def __init__(self, et):
+        self.elementtree = et
+        self.children = []
+        if et.docinfo.internalDTD:
+            self.children.append(Doctype(self, et.docinfo.root_name, 
+                                         et.docinfo.public_id, 
+                                         et.docinfo.system_url))
+        root = et.getroot()
+        node = root
+
+        while node.getprevious() is not None:
+            node = node.getprevious()
+        while node is not None:
+            self.children.append(node)
+            node = node.getnext()
+
+        self.text = None
+        self.tail = None
+    
+    def __getitem__(self, key):
+        return self.children[key]
+
+    def getnext(self):
+        return None
+
+class Doctype(object):
+    def __init__(self, root_node, name, public_id, system_id):
+        self.root_node = root_node
+        self.name = name
+        self.public_id = public_id
+        self.system_id = system_id
+        
+        self.text = None
+        self.tail = None
+
+    def getnext(self):
+        return self.root_node.children[1]
+
+class FragmentRoot(Root):
+    def __init__(self, children):
+        self.children = [FragmentWrapper(self, child) for child in children]
+        self.text = self.tail = None
+
+    def getnext(self):
+        return None
+
+class FragmentWrapper(object):
+    def __init__(self, fragment_root, obj):
+        self.root_node = fragment_root
+        self.obj = obj
+        if hasattr(self.obj, 'text'):
+            self.text = self.obj.text
+        else:
+            self.text = None
+        if hasattr(self.obj, 'tail'):
+            self.tail = self.obj.tail
+        else:
+            self.tail = None
+        self.isstring = isinstance(obj, basestring)
+        
+    def __getattr__(self, name):
+        return getattr(self.obj, name)
+    
+    def getnext(self):
+        siblings = self.root_node.children
+        idx = siblings.index(self)
+        if idx < len(siblings) - 1:
+            return siblings[idx + 1]
+        else:
+            return None
+
+    def __getitem__(self, key):
+        return self.obj[key]
+
+    def __nonzero__(self):
+        return bool(self.obj)
+
+    def getparent(self):
+        return None
+
+    def __str__(self):
+        return str(self.obj)
+
+        
 class TreeWalker(_base.NonRecursiveTreeWalker):
+    def __init__(self, tree):
+        if hasattr(tree, "getroot"):
+            tree = Root(tree)
+        elif isinstance(tree, list):
+            tree = FragmentRoot(tree)
+        _base.NonRecursiveTreeWalker.__init__(self, tree)
     def getNodeDetails(self, node):
         if isinstance(node, tuple): # Text node
             node, key = node
             assert key in ("text", "tail"), _("Text nodes are text or tail, found %s") % key
             return _base.TEXT, getattr(node, key)
 
-        if not(hasattr(node, "tag")):
-            node = node.getroot()
-
-        if node.tag in ("<DOCUMENT_ROOT>", "<DOCUMENT_FRAGMENT>"):
+        elif isinstance(node, Root):
             return (_base.DOCUMENT,)
 
-        elif node.tag == "<!DOCTYPE>":
-            return _base.DOCTYPE, node.text
+        elif isinstance(node, Doctype):
+            return _base.DOCTYPE, node.name, node.public_id, node.system_id
+
+        elif isinstance(node, FragmentWrapper) and node.isstring:
+            return _base.TEXT, node
 
         elif node.tag == etree.Comment:
             return _base.COMMENT, node.text
