feat(py): implement iter on ToNode (#1399)

CQCL · Aug 8, 2024 · e88910b · e88910b
1 parent 3609736
commit e88910b
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 1 deletion.
diff --git a/hugr-py/src/hugr/node_port.py b/hugr-py/src/hugr/node_port.py
@@ -90,6 +90,13 @@ def __getitem__(
     def out_port(self) -> OutPort:
         return OutPort(self.to_node(), 0)
 
+    def outputs(self) -> Iterator[OutPort]:
+        """Returns an iterator over the output ports of this node."""
+        return self[:]
+
+    def __iter__(self) -> Iterator[OutPort]:
+        return self.outputs()
+
     def inp(self, offset: PortOffset) -> InPort:
         """Generate an input port for this node.
 

diff --git a/hugr-py/tests/test_hugr_build.py b/hugr-py/tests/test_hugr_build.py
@@ -17,11 +17,14 @@
 def test_stable_indices():
     h = Hugr(ops.DFG([]))
 
-    nodes = [h.add_node(Not) for _ in range(3)]
+    nodes = [h.add_node(Not, num_outs=1) for _ in range(3)]
     assert len(h) == 4
     assert list(iter(h)) == [Node(i) for i in range(4)]
     assert all(data is not None for node, data in h.nodes())
 
+    assert len(list(nodes[0].outputs())) == 1
+    assert list(nodes[0]) == list(nodes[0].outputs())
+
     h.add_link(nodes[0].out(0), nodes[1].inp(0))
     assert h.children() == nodes