Fix panic for ancestors and descendants when the source node is invalid

releasenotes/notes/fix-ancestors-panic-3c64b9b43bb0551a.yaml

@@ -0,0 +1,5 @@
+fixes:
+  - |
+    Fixed a panic when passing an invalid source node to
+    :func:`~rustworkx.ancenstors` and :func:`~rustworkx.descendants`. See
+    `#1381 <https://github.com/Qiskit/rustworkx/issues/1381>`__ for more information.

src/traversal/mod.rs

@@ -30,7 +30,7 @@ use std::convert::TryFrom;
 
 use hashbrown::HashSet;
 
-use pyo3::exceptions::PyTypeError;
+use pyo3::exceptions::{PyIndexError, PyTypeError};
 use pyo3::prelude::*;
 use pyo3::Python;
 
@@ -219,11 +219,18 @@ pub fn bfs_predecessors(
 /// :rtype: set
 #[pyfunction]
 #[pyo3(text_signature = "(graph, node, /)")]
-pub fn ancestors(graph: &digraph::PyDiGraph, node: usize) -> HashSet<usize> {
-    core_ancestors(&graph.graph, NodeIndex::new(node))
+pub fn ancestors(graph: &digraph::PyDiGraph, node: usize) -> PyResult<HashSet<usize>> {
+    let index = NodeIndex::new(node);
+    if !graph.graph.contains_node(index) {
+        return Err(PyIndexError::new_err(format!(
+            "Node source index \"{}\" out of graph bound",
+            node
+        )));
+    }
+    Ok(core_ancestors(&graph.graph, index)
         .map(|x| x.index())
         .filter(|x| *x != node)
-        .collect()
+        .collect())
 }
 
 /// Return the descendants of a node in a graph.
@@ -240,12 +247,18 @@ pub fn ancestors(graph: &digraph::PyDiGraph, node: usize) -> HashSet<usize> {
 /// :rtype: set
 #[pyfunction]
 #[pyo3(text_signature = "(graph, node, /)")]
-pub fn descendants(graph: &digraph::PyDiGraph, node: usize) -> HashSet<usize> {
+pub fn descendants(graph: &digraph::PyDiGraph, node: usize) -> PyResult<HashSet<usize>> {
     let index = NodeIndex::new(node);
-    core_descendants(&graph.graph, index)
+    if !graph.graph.contains_node(index) {
+        return Err(PyIndexError::new_err(format!(
+            "Node source index \"{}\" out of graph bound",
+            node
+        )));
+    }
+    Ok(core_descendants(&graph.graph, index)
         .map(|x| x.index())
         .filter(|x| *x != node)
-        .collect()
+        .collect())
 }
 
 /// Breadth-first traversal of a directed graph with several source vertices.

tests/digraph/test_ancestors_descendants.py

@@ -39,6 +39,11 @@ def test_ancestors_no_descendants(self):
         res = rustworkx.ancestors(dag, node_b)
         self.assertEqual({node_a}, res)
 
+    def test_invalid_source(self):
+        graph = rustworkx.generators.directed_path_graph(5)
+        with self.assertRaises(IndexError):
+            rustworkx.ancestors(graph, 10)
+
 
 class TestDescendants(unittest.TestCase):
     def test_descendants(self):
@@ -62,3 +67,8 @@ def test_descendants_no_ancestors(self):
         node_c = dag.add_child(node_b, "c", {"b": 1})
         res = rustworkx.descendants(dag, node_b)
         self.assertEqual({node_c}, res)
+
+    def test_invalid_source(self):
+        graph = rustworkx.generators.directed_path_graph(5)
+        with self.assertRaises(IndexError):
+            rustworkx.descendants(graph, 10)