simplify interface for passing multiple tag selectors into element iteration methods

scoder · scoder · commit ba07885d1142 · 2012-04-24T15:23:31.000+02:00
diff --git a/doc/api.txt b/doc/api.txt
@@ -196,7 +196,7 @@ children.  Using the tree defined above, we get:
   >>> [ child.tag for child in root ]
   ['a', 'b', 'c', 'd']
 
-To iterate in the opposite direction, use the ``reversed()`` function
+To iterate in the opposite direction, use the builtin ``reversed()`` function
 that exists in Python 2.4 and later.
 
 Tree traversal should use the ``element.iter()`` method:
@@ -229,25 +229,33 @@ Note how ``element.iterdescendants()`` does not include the element
 itself, as opposed to ``element.iter()``.  The latter effectively
 implements the 'descendant-or-self' axis in XPath.
 
-All of these iterators support an additional ``tag`` keyword argument that
-filters the generated elements by tag name:
+All of these iterators support one (or more, since lxml 2.4) additional
+arguments that filter the generated elements by tag name:
 
 .. sourcecode:: pycon
 
-  >>> [ child.tag for child in root.iterchildren(tag='a') ]
+  >>> [ child.tag for child in root.iterchildren('a') ]
   ['a']
-  >>> [ child.tag for child in d.iterchildren(tag='a') ]
+  >>> [ child.tag for child in d.iterchildren('a') ]
   []
-  >>> [ el.tag for el in root.iterdescendants(tag='d') ]
+  >>> [ el.tag for el in root.iterdescendants('d') ]
   ['d']
-  >>> [ el.tag for el in root.iter(tag='d') ]
+  >>> [ el.tag for el in root.iter('d') ]
   ['d']
+  >>> [ el.tag for el in root.iter('d', 'a') ]
+  ['a', 'd']
+
+Note that the order of the elements is determined by the iteration order,
+which is the document order in most cases (except for preceding siblings
+and ancestors, where it is the reversed document order).  The order of
+the tag selection arguments is irrelevant, as you can see in the last
+example.
 
 The most common way to traverse an XML tree is depth-first, which
 traverses the tree in document order.  This is implemented by the
 ``.iter()`` method.  While there is no dedicated method for
 breadth-first traversal, it is almost as simple if you use the
-``collections.deque`` type from Python 2.4.
+``collections.deque`` type that is available in Python 2.4 and later.
 
 .. sourcecode:: pycon
 
diff --git a/doc/tutorial.txt b/doc/tutorial.txt
@@ -479,7 +479,8 @@ serialised the tree to XML:
     another - Child 3
 
 If you know you are only interested in a single tag, you can pass its name to
-``iter()`` to have it filter for you:
+``iter()`` to have it filter for you.  Since lxml 2.4, you can also pass more
+than one tag to intercept on multiple tags during iteration.
 
 .. sourcecode:: pycon
 
@@ -488,6 +489,12 @@ If you know you are only interested in a single tag, you can pass its name to
     child - Child 1
     child - Child 2
 
+    >>> for element in root.iter("another", "child"):
+    ...     print("%s - %s" % (element.tag, element.text))
+    child - Child 1
+    child - Child 2
+    another - Child 3
+
 By default, iteration yields all nodes in the tree, including
 ProcessingInstructions, Comments and Entity instances.  If you want to
 make sure only Element objects are returned, you can pass the
diff --git a/src/lxml/lxml.etree.pyx b/src/lxml/lxml.etree.pyx
@@ -1277,52 +1277,61 @@ cdef public class _Element [ type LxmlElementType, object LxmlElement ]:
             return None
         return _elementFactory(self._doc, c_node)
 
-    def itersiblings(self, tag=None, *, preceding=False):
-        u"""itersiblings(self, tag=None, preceding=False)
+    def itersiblings(self, tag=None, *tags, preceding=False):
+        u"""itersiblings(self, tag=None, *tags, preceding=False)
 
         Iterate over the following or preceding siblings of this element.
 
         The direction is determined by the 'preceding' keyword which
         defaults to False, i.e. forward iteration over the following
         siblings.  When True, the iterator yields the preceding
         siblings in reverse document order, i.e. starting right before
-        the current element and going left.  The generated elements
-        can be restricted to a specific tag name with the 'tag'
-        keyword.
+        the current element and going backwards.
+
+        The returned elements can be restricted to a specific tag name by
+        passing a tag or a series of tag names.
         """
-        return SiblingsIterator(self, tag, preceding=preceding)
+        if tag is not None:
+            tags += (tag,)
+        return SiblingsIterator(self, tags, preceding=preceding)
 
-    def iterancestors(self, tag=None):
-        u"""iterancestors(self, tag=None)
+    def iterancestors(self, tag=None, *tags):
+        u"""iterancestors(self, tag=None, *tags)
 
         Iterate over the ancestors of this element (from parent to parent).
 
-        The generated elements can be restricted to a specific tag name with
-        the 'tag' keyword.
+        The returned elements can be restricted to a specific tag name by
+        passing a tag or a series of tag names.
         """
-        return AncestorsIterator(self, tag)
+        if tag is not None:
+            tags += (tag,)
+        return AncestorsIterator(self, tags)
 
-    def iterdescendants(self, tag=None):
-        u"""iterdescendants(self, tag=None)
+    def iterdescendants(self, tag=None, *tags):
+        u"""iterdescendants(self, tag=None, *tags)
 
         Iterate over the descendants of this element in document order.
 
         As opposed to ``el.iter()``, this iterator does not yield the element
-        itself.  The generated elements can be restricted to a specific tag
-        name with the 'tag' keyword.
+        itself.  The returned elements can be restricted to a specific tag
+        name by passing a tag or a series of tag names.
         """
-        return ElementDepthFirstIterator(self, tag, inclusive=False)
+        if tag is not None:
+            tags += (tag,)
+        return ElementDepthFirstIterator(self, tags, inclusive=False)
 
-    def iterchildren(self, tag=None, *, reversed=False):
-        u"""iterchildren(self, tag=None, reversed=False)
+    def iterchildren(self, tag=None, *tags, reversed=False):
+        u"""iterchildren(self, tag=None, *tags, reversed=False)
 
         Iterate over the children of this element.
 
-        As opposed to using normal iteration on this element, the generated
-        elements can be restricted to a specific tag name with the 'tag'
-        keyword and reversed with the 'reversed' keyword.
+        As opposed to using normal iteration on this element, the returned
+        elements can be restricted to a specific tag name by passing a tag
+        or a series of tag names, and reversed with the 'reversed' keyword.
         """
-        return ElementChildIterator(self, tag, reversed=reversed)
+        if tag is not None:
+            tags += (tag,)
+        return ElementChildIterator(self, tags, reversed=reversed)
 
     def getroottree(self):
         u"""getroottree(self)
@@ -1336,15 +1345,17 @@ cdef public class _Element [ type LxmlElementType, object LxmlElement ]:
         _assertValidDoc(self._doc)
         return _elementTreeFactory(self._doc, None)
 
-    def getiterator(self, tag=None):
-        u"""getiterator(self, tag=None)
+    def getiterator(self, tag=None, *tags):
+        u"""getiterator(self, tag=None, *tags)
 
         Returns a sequence or iterator of all elements in the subtree in
         document order (depth first pre-order), starting with this
         element.
 
         Can be restricted to find only elements with a specific tag
-        (pass ``tag="xyz"``) or from a namespace (pass ``tag="{ns}*"``).
+        (pass ``"xyz"`` as tag) or from a namespace (pass ``"{ns}*"`` as tag).
+        Passing a sequence of tags will let the iterator return all
+        elements matching any of these tags, in document order.
 
         You can also pass the Element, Comment, ProcessingInstruction and
         Entity factory functions to look only for the specific element type.
@@ -1357,34 +1368,43 @@ cdef public class _Element [ type LxmlElementType, object LxmlElement ]:
           method in new code if you require backwards compatibility
           with older versions of lxml or ElementTree.
         """
-        return ElementDepthFirstIterator(self, tag)
+        if tag is not None:
+            tags += (tag,)
+        return ElementDepthFirstIterator(self, tags)
 
-    def iter(self, tag=None):
-        u"""iter(self, tag=None)
+    def iter(self, tag=None, *tags):
+        u"""iter(self, tag=None, *tags)
 
         Iterate over all elements in the subtree in document order (depth
         first pre-order), starting with this element.
 
         Can be restricted to find only elements with a specific tag
-        (pass ``tag="xyz"``) or from a namespace (pass ``tag="{ns}*"``).
+        (pass ``"xyz"`` as tag) or from a namespace (pass ``"{ns}*"`` as tag).
+        Passing a sequence of tags will let the iterator return all
+        elements matching any of these tags, in document order.
 
         You can also pass the Element, Comment, ProcessingInstruction and
         Entity factory functions to look only for the specific element type.
         """
-        return ElementDepthFirstIterator(self, tag)
+        if tag is not None:
+            tags += (tag,)
+        return ElementDepthFirstIterator(self, tags)
 
-    def itertext(self, tag=None, *, with_tail=True):
-        u"""itertext(self, tag=None, with_tail=True)
+    def itertext(self, tag=None, *tags, with_tail=True):
+        u"""itertext(self, tag=None, *tags, with_tail=True)
 
         Iterates over the text content of a subtree.
 
-        You can pass the ``tag`` keyword argument to restrict text content to
-        a specific tag name.
+        You can pass a tag name to restrict text content to a specific tag
+        name.  Passing a sequence of tags will let the iterator consider
+        all elements matching any of these tags.
 
         You can set the ``with_tail`` keyword argument to ``False`` to skip
         over tail text.
         """
-        return ElementTextIterator(self, tag, with_tail=with_tail)
+        if tag is not None:
+            tags += (tag,)
+        return ElementTextIterator(self, tags, with_tail=with_tail)
 
     def makeelement(self, _tag, attrib=None, nsmap=None, **_extra):
         u"""makeelement(self, _tag, attrib=None, nsmap=None, **_extra)
@@ -1915,15 +1935,16 @@ cdef public class _ElementTree [ type LxmlElementTreeType,
         tree.xmlFree(c_path)
         return path
 
-    def getiterator(self, tag=None):
-        u"""getiterator(self, tag=None)
+    def getiterator(self, tag=None, *tags):
+        u"""getiterator(self, *tags, tag=None)
 
         Returns a sequence or iterator of all elements in document order
         (depth first pre-order), starting with the root element.
 
         Can be restricted to find only elements with a specific tag
-        (pass ``tag="xyz"`` or ``tag="{ns}xyz"``) or from a namespace
-        (pass ``tag="{ns}*"``).
+        (pass ``"xyz"`` as tag) or from a namespace (pass ``"{ns}*"`` as tag).
+        Passing a sequence of tags will let the iterator return all
+        elements matching any of these tags, in document order.
 
         You can also pass the Element, Comment, ProcessingInstruction and
         Entity factory functions to look only for the specific element type.
@@ -1939,18 +1960,22 @@ cdef public class _ElementTree [ type LxmlElementTreeType,
         root = self.getroot()
         if root is None:
             return ()
-        return root.getiterator(tag)
+        if tag is not None:
+            tags += (tag,)
+        return root.getiterator(*tags)
 
-    def iter(self, tag=None):
-        u"""iter(self, tag=None)
+    def iter(self, tag=None, *tags):
+        u"""iter(self, tag=None, *tags)
 
         Creates an iterator for the root element.  The iterator loops over
         all elements in this tree, in document order.
         """
         root = self.getroot()
         if root is None:
             return ()
-        return root.iter(tag)
+        if tag is not None:
+            tags += (tag,)
+        return root.iter(*tags)
 
     def find(self, path, namespaces=None):
         u"""find(self, path, namespaces=None)
@@ -2424,24 +2449,24 @@ cdef class _MultiTagMatcher:
         self._tag_count = 0
         if self._cached_tags:
             for i in xrange(count):
-                python.Py_XDECREF(self._cached_tags[i].href)
+                cpython.ref.Py_XDECREF(self._cached_tags[i].href)
             cpython.mem.PyMem_Free(self._cached_tags)
             self._cached_tags = NULL
 
     cdef initTagMatch(self, tags):
         self._cached_doc = None
         del self._py_tags[:]
         self._clear()
-        if tags is None:
-            # match anything
+        if tags is None or tags == ():
+            # no selection in tags argument => match anything
             self._node_types = (
                 1 << tree.XML_COMMENT_NODE |
                 1 << tree.XML_PI_NODE |
                 1 << tree.XML_ENTITY_REF_NODE |
                 1 << tree.XML_ELEMENT_NODE)
-            return
-        self._node_types = 0
-        self._storeTags(tags, set())
+        else:
+            self._node_types = 0
+            self._storeTags(tags, set())
 
     cdef _storeTags(self, tag, set seen):
         if tag is Comment:
@@ -2456,10 +2481,13 @@ cdef class _MultiTagMatcher:
             if tag in seen:
                 return
             seen.add(tag)
-            href, name = _getNsTag(tag)
-            if name == b'*':
-                name = None
-            self._py_tags.append((href, name))
+            if tag in ('*', '{*}*'):
+                self._node_types |= 1 << tree.XML_ELEMENT_NODE
+            else:
+                href, name = _getNsTag(tag)
+                if name == b'*':
+                    name = None
+                self._py_tags.append((href, name))
         else:
             # support a sequence of tags
             for item in tag:
@@ -2476,6 +2504,7 @@ cdef class _MultiTagMatcher:
         self._tag_count = 0
         if not self._py_tags:
             self._cached_doc = doc
+            self._cached_size = dict_size
             return 0
         if not self._cached_tags:
             self._cached_tags = <qname*>cpython.mem.PyMem_Malloc(len(self._py_tags) * sizeof(qname))
@@ -2599,7 +2628,7 @@ cdef class ElementDepthFirstIterator:
     instructions.  To filter them out, check if the ``tag`` property
     of the returned element is a string (i.e. not None and not a
     factory function), or pass the ``Element`` factory for the ``tag``
-    keyword.
+    argument to receive only Elements.
 
     If the optional ``tag`` argument is not None, the iterator returns only
     the elements that match the respective name and namespace.
diff --git a/src/lxml/tests/test_etree.py b/src/lxml/tests/test_etree.py
@@ -2195,19 +2195,46 @@ def test_getiterator_filter_multiple(self):
 
         self.assertEquals(
             [a, b],
-               list(a.getiterator( ('a', 'b') )))
+               list(a.getiterator('a', 'b')))
+        self.assertEquals(
+            [],
+              list(a.getiterator('x', 'y')))
+        self.assertEquals(
+            [a, f],
+              list(a.getiterator('f', 'a')))
+        self.assertEquals(
+            [c, e, f],
+               list(c.getiterator('c', '*', 'a')))
+        self.assertEquals(
+            [],
+                  list(a.getiterator( (), () )))
+
+    def test_getiterator_filter_multiple_tuple(self):
+        Element = self.etree.Element
+        SubElement = self.etree.SubElement
+
+        a = Element('a')
+        b = SubElement(a, 'b')
+        c = SubElement(a, 'c')
+        d = SubElement(b, 'd')
+        e = SubElement(c, 'e')
+        f = SubElement(c, 'f')
+
+        self.assertEquals(
+            [a, b],
+                  list(a.getiterator( ('a', 'b') )))
         self.assertEquals(
             [],
               list(a.getiterator( ('x', 'y') )))
         self.assertEquals(
             [a, f],
-              list(a.getiterator( ('f', 'a') )))
+                  list(a.getiterator( ('f', 'a') )))
         self.assertEquals(
             [c, e, f],
-               list(c.getiterator( ('c', '*', 'a') )))
+                     list(c.getiterator( ('c', '*', 'a') )))
         self.assertEquals(
             [],
-                  list(a.getiterator( () )))
+              list(a.getiterator( () )))
 
     def test_getiterator_filter_namespace(self):
         Element = self.etree.Element
