fix get_w_partner util function for multigraph

zxcalc · Jul 9, 2024 · 6126178 · 6126178
1 parent 716c3e9
commit 6126178
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/pyzx/utils.py b/pyzx/utils.py
@@ -60,8 +60,9 @@ def vertex_is_w(ty: VertexType) -> bool:
 
 def get_w_partner(g, v):
     assert vertex_is_w(g.type(v))
-    for u in g.neighbors(v):
-        if g.edge_type((u, v)) == EdgeType.W_IO:
+    for edge in g.incident_edges(v):
+        if g.edge_type(edge) == EdgeType.W_IO:
+            u = edge[0] if edge[1] == v else edge[1]
             assert vertex_is_w(g.type(u))
             return u
     assert False