From e88910bc50f6459213e17d0ce5a8c1b87037b9a6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Agust=C3=ADn=20Borgna?=
 <121866228+aborgna-q@users.noreply.github.com>
Date: Thu, 8 Aug 2024 16:27:24 +0100
Subject: [PATCH] feat(py): implement `iter` on `ToNode` (#1399)

---
 hugr-py/src/hugr/node_port.py    | 7 +++++++
 hugr-py/tests/test_hugr_build.py | 5 ++++-
 2 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/hugr-py/src/hugr/node_port.py b/hugr-py/src/hugr/node_port.py
index 04f0804ab..66b9e8ce2 100644
--- 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
index de7a1646c..021cfc5cf 100644
--- 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