Add method: has_node

clippy Add tests
Qiskit · Apr 22, 2024 · 70b480b · 70b480b
1 parent ab93329
commit 70b480b
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 6 deletions.
diff --git a/src/digraph.rs b/src/digraph.rs
@@ -595,6 +595,18 @@ impl PyDiGraph {
         self.node_indices()
     }
 
+    /// Return True if there is a node in the graph.
+    ///
+    /// :param int node: The node index to check
+    ///
+    /// :returns: True if there is a node false if there is no node
+    /// :rtype: bool
+    #[pyo3(text_signature = "(self, node, /)")]
+    pub fn has_node(&self, node: usize) -> bool {
+        let index = NodeIndex::new(node);
+        return self.graph.node_weight(index).is_some();
+    }
+
     /// Return True if there is an edge from node_a to node_b.
     ///
     /// :param int node_a: The source node index to check for an edge
@@ -3059,7 +3071,7 @@ impl PyDiGraph {
     /// required to return a boolean value stating whether the node's data payload fits some criteria.
     ///
     /// For example::
-    ///     
+    ///
     ///     from rustworkx import PyDiGraph
     ///
     ///     graph = PyDiGraph()
@@ -3107,8 +3119,8 @@ impl PyDiGraph {
     ///     def my_filter_function(edge):
     ///         if edge:
     ///             return edge == 'B'
-    ///         return False  
-    ///        
+    ///         return False
+    ///
     ///     indices = graph.filter_edges(my_filter_function)
     ///     assert indices == [1]
     ///

diff --git a/src/graph.rs b/src/graph.rs
@@ -468,6 +468,18 @@ impl PyGraph {
         self.node_indices()
     }
 
+    /// Return True if there is a node.
+    ///
+    /// :param int node: The index for the node
+    ///
+    /// :returns: True if there is a node false if there is no node
+    /// :rtype: bool
+    #[pyo3(text_signature = "(self, node, /)")]
+    pub fn has_node(&self, node: usize) -> bool {
+        let index = NodeIndex::new(node);
+        return self.graph.node_weight(index).is_some();
+    }
+
     /// Return True if there is an edge between ``node_a`` and ``node_b``.
     ///
     /// :param int node_a: The index for the first node
@@ -2039,7 +2051,7 @@ impl PyGraph {
     /// required to return a boolean value stating whether the node's data payload fits some criteria.
     ///
     /// For example::
-    ///     
+    ///
     ///     from rustworkx import PyGraph
     ///
     ///     graph = PyGraph()
@@ -2087,8 +2099,8 @@ impl PyGraph {
     ///     def my_filter_function(edge):
     ///         if edge:
     ///             return edge == 'B'
-    ///         return False  
-    ///        
+    ///         return False
+    ///
     ///     indices = graph.filter_edges(my_filter_function)
     ///     assert indices == [1]
     ///

diff --git a/tests/digraph/test_nodes.py b/tests/digraph/test_nodes.py
@@ -85,6 +85,7 @@ def test_remove_nodes_retain_edges_single_edge(self):
         res = dag.nodes()
         self.assertEqual(["a", "c"], res)
         self.assertEqual([0, 2], dag.node_indexes())
+        self.assertTrue(dag.has_node(node_a))
         self.assertTrue(dag.has_edge(node_a, node_c))
         self.assertEqual(dag.get_all_edge_data(node_a, node_c), ["Edgy"])
 
@@ -97,6 +98,7 @@ def test_remove_nodes_retain_edges_single_edge_outgoing_weight(self):
         res = dag.nodes()
         self.assertEqual(["a", "c"], res)
         self.assertEqual([0, 2], dag.node_indexes())
+        self.assertTrue(dag.has_node(node_a))
         self.assertTrue(dag.has_edge(node_a, node_c))
         self.assertEqual(dag.get_all_edge_data(node_a, node_c), ["Edgy_mk2"])