Add in and out edge indices functions (#1369)

* Add method to retrieve incoming edge indices for a node

* Add type annotations in stub file

* Add tests for in_edge_indices

* Add method to get outgoing edge indices for a node

* Add in_edge_indices and out_edge_indices methods in PyGraph

* update stub file

* Add tests for out_edge_indices and  in_edge_indices

* Add release notes for in_edge_indices and out_edge_indices functions

* Change releasenote name

Author: ThisuraGallage

releasenotes/notes/add-in-out-edge-indices-functions-fc1ce122ff761ad6.yaml

@@ -0,0 +1,7 @@
+features:
+  - |
+    Added two new functions, :func:`~rustworkx.in_edge_indices` and
+    :func:`~rustworkx.out_edge_indices`, to the ``rustworkx.PyDiGraph`` class. These functions return the indices of all incoming and outgoing edges for a given node in a directed graph, respectively.
+  - |
+    Added two new functions, :func:`~rustworkx.in_edge_indices` and
+    :func:`~rustworkx.out_edge_indices`, to the ``rustworkx.PyGraph`` class. As ``PyGraph`` is an undirected graph, both functions return the indices of all edges connected to the given node. This maintains API consistency with the directed graph implementation.

rustworkx/rustworkx.pyi

@@ -1262,13 +1262,15 @@ class PyGraph(Generic[_S, _T]):
     def in_edges(self, node: int, /) -> WeightedEdgeList[_T]: ...
     def incident_edge_index_map(self, node: int, /) -> EdgeIndexMap: ...
     def incident_edges(self, node: int, /) -> EdgeIndices: ...
+    def in_edge_indices(self, node: int, /) -> EdgeIndices: ...
     def neighbors(self, node: int, /) -> NodeIndices: ...
     def node_indexes(self) -> NodeIndices: ...
     def node_indices(self) -> NodeIndices: ...
     def nodes(self) -> list[_S]: ...
     def num_edges(self) -> int: ...
     def num_nodes(self) -> int: ...
     def out_edges(self, node: int, /) -> WeightedEdgeList[_T]: ...
+    def out_edge_indices(self, node: int, /) -> EdgeIndices: ...
     @staticmethod
     def read_edge_list(
         path: str,
@@ -1428,6 +1430,7 @@ class PyDiGraph(Generic[_S, _T]):
     def in_edges(self, node: int, /) -> WeightedEdgeList[_T]: ...
     def incident_edge_index_map(self, node: int, /, all_edges: bool = ...) -> EdgeIndexMap: ...
     def incident_edges(self, node: int, /, all_edges: bool = ...) -> EdgeIndices: ...
+    def in_edge_indices(self, node: int, /) -> EdgeIndices: ...
     def insert_node_on_in_edges(self, node: int, ref_node: int, /) -> None: ...
     def insert_node_on_in_edges_multiple(self, node: int, ref_nodes: Sequence[int], /) -> None: ...
     def insert_node_on_out_edges(self, node: int, ref_node: int, /) -> None: ...
@@ -1444,6 +1447,7 @@ class PyDiGraph(Generic[_S, _T]):
     def num_nodes(self) -> int: ...
     def out_degree(self, node: int, /) -> int: ...
     def out_edges(self, node: int, /) -> WeightedEdgeList[_T]: ...
+    def out_edge_indices(self, node: int, /) -> EdgeIndices: ...
     def predecessor_indices(self, node: int, /) -> NodeIndices: ...
     def predecessors(self, node: int, /) -> list[_S]: ...
     @staticmethod

src/digraph.rs

@@ -1842,6 +1842,54 @@ impl PyDiGraph {
         }
     }
 
+    /// Return the list of incoming edge indices to a provided node
+    ///
+    /// This method will return the incoming edges of the provided
+    /// ``node``.
+    ///
+    /// :param int node: The node index to get incoming edges from. If
+    ///     this node index is not present in the graph this method will
+    ///     return an empty list and not error.
+    ///
+    /// :returns: A list of the incoming edge indices to a node in the graph
+    /// :rtype: EdgeIndices
+    #[pyo3(text_signature = "(self, node, /)")]
+    pub fn in_edge_indices(&self, node: usize) -> EdgeIndices {
+        let node_index = NodeIndex::new(node);
+        let dir = petgraph::Direction::Incoming;
+        EdgeIndices {
+            edges: self
+                .graph
+                .edges_directed(node_index, dir)
+                .map(|e| e.id().index())
+                .collect(),
+        }
+    }
+
+    /// Return the list of outgoing edge indices from a provided node
+    ///
+    /// This method will return the outgoing edges of the provided
+    /// ``node``.
+    ///
+    /// :param int node: The node index to get outgoing edges from. If
+    ///     this node index is not present in the graph this method will
+    ///     return an empty list and not error.
+    ///
+    /// :returns: A list of the outgoing edge indices from a node in the graph
+    /// :rtype: EdgeIndices
+    #[pyo3(text_signature = "(self, node, /)")]
+    pub fn out_edge_indices(&self, node: usize) -> EdgeIndices {
+        let node_index = NodeIndex::new(node);
+        let dir = petgraph::Direction::Outgoing;
+        EdgeIndices {
+            edges: self
+                .graph
+                .edges_directed(node_index, dir)
+                .map(|e| e.id().index())
+                .collect(),
+        }
+    }
+
     /// Return the index map of edges incident to a provided node
     ///
     /// By default this method will only return the outgoing edges of

src/graph.rs

@@ -552,6 +552,52 @@ impl PyGraph {
         }
     }
 
+    /// Return the list of edge indices incident to a provided node.
+    ///
+    /// This method returns the indices of all edges connected to the provided
+    /// ``node``. In undirected graphs, all edges connected to the node are
+    /// returned as there is no distinction between incoming and outgoing edges.
+    ///
+    /// :param int node: The node index to get incident edges from. If
+    ///     this node index is not present in the graph this method will
+    ///     return an empty list and not error.
+    ///
+    /// :returns: A list of the edge indices incident to the node
+    /// :rtype: EdgeIndices
+    #[pyo3(text_signature = "(self, node, /)")]
+    pub fn in_edge_indices(&self, node: usize) -> EdgeIndices {
+        EdgeIndices {
+            edges: self
+                .graph
+                .edges(NodeIndex::new(node))
+                .map(|e| e.id().index())
+                .collect(),
+        }
+    }
+
+    /// Return the list of edge indices incident to a provided node.
+    ///
+    /// This method returns the indices of all edges connected to the provided
+    /// ``node``. In undirected graphs, all edges connected to the node are
+    /// returned as there is no distinction between incoming and outgoing edges.
+    ///
+    /// :param int node: The node index to get incident edges from. If
+    ///     this node index is not present in the graph this method will
+    ///     return an empty list and not error.
+    ///
+    /// :returns: A list of the edge indices incident to the node
+    /// :rtype: EdgeIndices
+    #[pyo3(text_signature = "(self, node, /)")]
+    pub fn out_edge_indices(&self, node: usize) -> EdgeIndices {
+        EdgeIndices {
+            edges: self
+                .graph
+                .edges(NodeIndex::new(node))
+                .map(|e| e.id().index())
+                .collect(),
+        }
+    }
+
     /// Return the index map of edges incident to a provided node
     ///
     /// :param int node: The node index to get incident edges from. If

tests/digraph/test_edges.py

@@ -813,6 +813,40 @@ def test_incident_edges_all_edges(self):
         res = graph.incident_edges(node_d, all_edges=True)
         self.assertEqual([2, 1], res)
 
+    def test_in_edge_indices(self):
+        graph = rustworkx.PyDiGraph()
+        node_a = graph.add_node(0)
+        node_b = graph.add_node(1)
+        node_c = graph.add_node("c")
+        node_d = graph.add_node("d")
+        edge_ac = graph.add_edge(node_a, node_c, "edge a")
+        graph.add_edge(node_b, node_d, "edge b")
+        edge_dc = graph.add_edge(node_d, node_c, "edge c")
+        res = graph.in_edge_indices(node_c)
+        self.assertEqual([edge_dc, edge_ac], res)
+
+    def test_in_edge_indices_invalid_node(self):
+        graph = rustworkx.PyDiGraph()
+        res = graph.in_edge_indices(0)
+        self.assertEqual([], res)
+
+    def test_out_edge_indices(self):
+        graph = rustworkx.PyDiGraph()
+        node_a = graph.add_node(0)
+        node_b = graph.add_node(1)
+        node_c = graph.add_node("c")
+        node_d = graph.add_node("d")
+        edge_ab = graph.add_edge(node_a, node_b, "edge a")
+        edge_ac = graph.add_edge(node_a, node_c, "edge b")
+        graph.add_edge(node_c, node_d, "edge c")
+        res = graph.out_edge_indices(node_a)
+        self.assertEqual([edge_ac, edge_ab], res)
+
+    def test_out_edge_indices_invalid_node(self):
+        graph = rustworkx.PyDiGraph()
+        res = graph.out_edge_indices(0)
+        self.assertEqual([], res)
+
     def test_incident_edge_index_map(self):
         graph = rustworkx.PyDiGraph()
         node_a = graph.add_node(0)

tests/graph/test_edges.py

@@ -526,6 +526,40 @@ def test_incident_edges_invalid_node(self):
         res = graph.incident_edges(42)
         self.assertEqual([], res)
 
+    def test_in_edge_indices(self):
+        graph = rustworkx.PyGraph()
+        node_a = graph.add_node(0)
+        node_b = graph.add_node(1)
+        node_c = graph.add_node("c")
+        node_d = graph.add_node("d")
+        graph.add_edge(node_a, node_c, "edge a")
+        graph.add_edge(node_b, node_d, "edge_b")
+        graph.add_edge(node_c, node_d, "edge c")
+        res = graph.in_edge_indices(node_d)
+        self.assertEqual({1, 2}, set(res))
+
+    def test_in_edge_indices_invalid_node(self):
+        graph = rustworkx.PyGraph()
+        res = graph.in_edge_indices(0)
+        self.assertEqual([], res)
+
+    def test_out_edge_indices(self):
+        graph = rustworkx.PyGraph()
+        node_a = graph.add_node(0)
+        node_b = graph.add_node(1)
+        node_c = graph.add_node("c")
+        node_d = graph.add_node("d")
+        graph.add_edge(node_a, node_c, "edge a")
+        graph.add_edge(node_b, node_d, "edge_b")
+        graph.add_edge(node_c, node_d, "edge c")
+        res = graph.out_edge_indices(node_d)
+        self.assertEqual({1, 2}, set(res))
+
+    def test_out_edge_indices_invalid_node(self):
+        graph = rustworkx.PyGraph()
+        res = graph.out_edge_indices(0)
+        self.assertEqual([], res)
+
     def test_incident_edge_index_map(self):
         graph = rustworkx.PyGraph()
         node_a = graph.add_node(0)