Use iterators where possible.

STIXProject · May 11, 2015 · 204b9cb · 204b9cb
1 parent 501318e
commit 204b9cb
Show file tree

Hide file tree

Showing 4 changed files with 24 additions and 11 deletions.
diff --git a/ramrod/base.py b/ramrod/base.py
@@ -394,7 +394,7 @@ def _is_empty(cls, node):
         attributes or text values.
 
         """
-        nodes = node.xpath(xmlconst.XPATH_DESCENDANT_OR_SELF)
+        nodes = node.iter('*')
         content = any(x.attrib or utils.strip_whitespace(x.text) for x in nodes)
         return content is False
 
@@ -508,7 +508,7 @@ def _get_duplicates(self, root):
         namespaces = self.NSMAP.values()
         id_nodes = collections.defaultdict(list)
 
-        for desc in utils.descendants(root):
+        for desc in utils.iterdescendants(root):
             if 'id' not in desc.attrib:
                 continue
 
@@ -678,7 +678,7 @@ def _apply_namespace_updates(self, root):
         a new `ns0` namespace alias.
 
         """
-        for node in utils.descendants(root):
+        for node in utils.iterdescendants(root):
             node_ns = utils.get_namespace(node)
             updated_ns = self.UPDATE_NS_MAP.get(node_ns, node_ns)
             node.tag = node.tag.replace(node_ns, updated_ns)

diff --git a/ramrod/stix/stix_1_0.py b/ramrod/stix/stix_1_0.py
@@ -58,7 +58,7 @@ def _check_maec(cls, node):
 
         """
         try:
-            namespaces = (utils.get_ext_namespace(x) for x in utils.children(node))
+            namespaces = (utils.get_ext_namespace(x) for x in utils.iterchildren(node))
             return all(ns == cls.NS_MAEC_EXT for ns in namespaces)
         except KeyError:
             # At least one node didn't contain an xsi:type attribute
@@ -96,7 +96,7 @@ def _check_capec(cls, node):
 
         """
         try:
-            namespaces = (utils.get_ext_namespace(x) for x in utils.children(node))
+            namespaces = (utils.get_ext_namespace(x) for x in utils.iterchildren(node))
             return all(ns == cls.NS_CAPEC_EXT for ns in namespaces)
         except KeyError:
             # At least one node didn't contain an xsi:type attribute

diff --git a/ramrod/stix/stix_1_0_1.py b/ramrod/stix/stix_1_0_1.py
@@ -68,7 +68,7 @@ def _check_maec(cls, node):
         """
         try:
 
-            namespaces = (utils.get_ext_namespace(x) for x in utils.children(node))
+            namespaces = (utils.get_ext_namespace(x) for x in utils.iterchildren(node))
             return all(ns == cls.NS_MAEC_EXT for ns in namespaces)
         except KeyError:
             # At least one node didn't contain an xsi:type attribute
@@ -105,7 +105,7 @@ def _check_capec(cls, node):
 
         """
         try:
-            namespaces = (utils.get_ext_namespace(x) for x in utils.children(node))
+            namespaces = (utils.get_ext_namespace(x) for x in utils.iterchildren(node))
             return all(ns == cls.NS_CAPEC_EXT for ns in namespaces)
         except KeyError:
             # At least one node didn't contain an xsi:type attribute

diff --git a/ramrod/utils.py b/ramrod/utils.py
@@ -305,20 +305,33 @@ def validate_versions(from_, to_, allowed):
         raise errors.InvalidVersionError(error)
 
 
+def iterchildren(node):
+    """Returns an iterator which yields direct child elements of `node`.
+
+    """
+    return node.iterchildren('*')
+
+
 def children(node):
     """Returns an iterable collection of etree Element nodes that are direct
     children of `node`.
 
     """
-    return node.xpath(xmlconst.XPATH_RELATIVE_CHILDREN)
+    return list(iterchildren(node))
+
+
+def iterdescendants(node):
+    """Returns an iterator which yields descendant elements of `node`.
+
+    """
+    return node.iterdescendants('*')
 
 
 def descendants(node):
-    """Returns an iterable collection of etree Element nodes that are
-    descendants of `node`.
+    """Returns a list of etree Element nodes that are descendants of `node`.
 
     """
-    return node.xpath(xmlconst.XPATH_RELATIVE_DESCENDANTS)
+    return list(iterdescendants(node))
 
 
 def strip_whitespace(string):