core: add iteration to :class:Node

Iteration can be in DFS/BFS order, with/out duplicates, and on
Nodes/Edges.

Signed-off-by: Amit Beka <[email protected]>

Author: Amit Beka

core.py

@@ -523,6 +523,49 @@ def equals(self, other, *, recursive=True, ordered=False):
             return False
         return True
 
+    def iter(self, obj="nodes", method="dfs", duplicates=False, key=None):
+        """Iterates the :class:Node objects in the subtree of self.
+
+        Args:
+            obj: yield Node objects (use value "nodes", default) or Edge
+                objects (use values "edges")
+            method: do breadth-first iteration (use value "bfs") or depth-dirst
+                iteration (value "dfs", default).
+            duplicates: If True, may return the same object twice if it is
+                encountered twice, because of the DAG structure which isn't
+                necessarily a tree. If it is False, all objects will be yielded
+                only the first time they are encountered. Defaults to False.
+            key: boolean function that filters the iterable items. key function
+                takes one argument (the item) and returns True if it should be
+                returned to the user. If an item isn't returned, its subtree
+                is still iterated.  Defaults to None (returns all items).
+
+        Yields:
+            a :class:Node or :class:Edge object according to the iteration
+            parameters.
+
+        """
+        if method not in ("dfs", "bfs"):
+            raise ValueError("method can be either 'dfs' or 'bfs'")
+        if obj not in ("nodes", "edges"):
+            raise ValueError("obj can be either 'nodes' or 'edges'")
+        processed = set()
+        if obj == 'nodes':
+            waiting = [self]
+        else:
+            waiting = self._outgoing[:]
+        while len(waiting):
+            curr = waiting.pop(0)
+            if key is None or key(curr):
+                yield curr
+            processed.add(curr)
+            to_add = curr.children if obj == 'nodes' else list(curr.child)
+            to_add = [x for x in to_add if duplicates or x not in processed]
+            if method == "bfs":
+                waiting.extend(to_add)
+            else:
+                waiting = to_add + waiting
+
 
 class Layer:
     """Group of similar :class:Node objects in UCCA annotation graph.

tests/test_ucca_ut.py

@@ -204,6 +204,28 @@ def test_copying(self):
         p2 = p1.copy([l0id])
         self.assertTrue(p1.layer(l0id).equals(p2.layer(l0id)))
 
+    def test_iteration(self):
+        p = self._create_basic_passage()
+        l1, l2 = p.layer('1'), p.layer('2')
+        node11, node12, node13 = l1.all
+        node22, node21 = l2.all
+
+        self.assertSequenceEqual(list(node11.iter()), [node11])
+        self.assertSequenceEqual(list(node11.iter(obj='edges')), [])
+        self.assertSequenceEqual(list(node13.iter(key=lambda x: x.tag != '3')),
+                                 [])
+        self.assertSequenceEqual(list(node12.iter()), [node12, node13, node11])
+        self.assertSequenceEqual(list(x.ID for x in node12.iter(obj='edges')),
+                                 ['1.2->1.3', '1.2->1.1'])
+        self.assertSequenceEqual(list(node21.iter(duplicates=True)),
+                                 [node21, node11, node12, node13, node11])
+        self.assertSequenceEqual(list(node21.iter()),
+                                 [node21, node11, node12, node13])
+        self.assertSequenceEqual(list(node22.iter(method='bfs',
+                                                  duplicates=True)),
+                                 [node22, node11, node12, node13, node13,
+                                  node11])
+
 
 class Layer0Tests(unittest.TestCase):
     """Tests module layer0 functionality."""