Commits

Anonymous committed f7db63b

Tree - Add element only tree iterator

Comments (0)

Files changed (2)

emeraldtree/tests/test_tree.py

     elem.append(u'b')
     assert len(list(elem.iter_elements())) == 1
 
+def test_Element_iter_elements_tree():
+    elem = Element(u'a')
+    assert len(list(elem.iter_elements_tree())) == 1
+
+    elem.append(u'b')
+    assert len(list(elem.iter_elements_tree())) == 1
+
+    elem.append(Element(u'c'))
+    assert len(list(elem.iter_elements_tree())) == 2
+
+    elem.append(u'd')
+    assert len(list(elem.iter_elements_tree())) == 2
+
 def test_Comment():
     elem = Comment(u'a')
     assert serialize(elem) == u'<!--a-->'

emeraldtree/tree.py

             if child.__class__ is Element:
                 yield child
 
+    def iter_elements_tree(self):
+        """
+        Creates an interator over all elements in document order.
+        """
+        work = [(self ,)]
+        while work:
+            cur = work.pop()
+            for i in cur:
+                yield i
+                work.append(i.iter_elements())
+
 
 ##
 # Subelement factory.  This function creates an element instance, and
                 self._raise_serialization_error(qname)
 
         # populate qname and namespaces table
-        if isinstance(elem, Element):
-            for elem in elem.iter():
-                if isinstance(elem, Element):
-                    tag = elem.tag
-                    if isinstance(tag, QName):
-                        add_qname(tag)
-                    elif isinstance(tag, basestring):
-                        add_qname(QName(tag))
-                    elif tag is not None:
-                        self._raise_serialization_error(tag)
+        if elem.__class__ is Element:
+            for elem in elem.iter_elements_tree():
+                tag = elem.tag
+                if isinstance(tag, QName):
+                    add_qname(tag)
+                elif isinstance(tag, basestring):
+                    add_qname(QName(tag))
+                elif tag is not None:
+                    self._raise_serialization_error(tag)
 
-                    for key in elem.keys():
-                        if isinstance(key, QName):
-                            add_qname(key)
-                        elif isinstance(key, basestring):
-                            add_qname(QName(key))
-                        elif key is not None:
-                            self._raise_serialization_error(key)
+                for key in elem.attrib.iterkeys():
+                    if isinstance(key, QName):
+                        add_qname(key)
+                    elif isinstance(key, basestring):
+                        add_qname(QName(key))
+                    elif key is not None:
+                        self._raise_serialization_error(key)
 
         return qnames, used_namespaces