Node and Edge Filtering

releasenotes/notes/add-filter-98d00f306b5689ee.yaml

@@ -0,0 +1,43 @@
+---
+features:
+  - |
+    The :class:`~rustworkx.PyGraph` and the :class:`~rustworkx.PyDiGraph` classes have a new method 
+    :meth:`~rustworkx.PyGraph.filter_nodes` (or :meth:`~rustworkx.PyDiGraph.filter_nodes`). 
+    This method returns a :class:`~.NodeIndices` object with the resulting nodes that fit some abstract criteria indicated by a filter function.
+    For example:
+
+    .. jupyter-execute::
+
+        from rustworkx import PyGraph
+
+        graph = PyGraph()
+        graph.add_nodes_from(list(range(5))) # Adds nodes from 0 to 5
+
+        def my_filter_function(node):
+          return node > 2
+
+        indices = graph.filter_nodes(my_filter_function)
+        print(indices)
+
+  - |
+    The :class:`~rustworkx.PyGraph` and the :class:`~rustworkx.PyDiGraph` classes have a new method 
+    :meth:`~rustworkx.PyGraph.filter_edges` (or :meth:`~rustworkx.PyDiGraph.filter_edges`). 
+    This method returns a :class:`~.EdgeIndices` object with the resulting edges that fit some abstract criteria indicated by a filter function.
+    For example:
+
+    .. jupyter-execute::
+
+        from rustworkx import PyGraph
+        from rustworkx.generators import complete_graph
+
+        graph = PyGraph()
+        graph.add_nodes_from(range(3))
+        graph.add_edges_from([(0, 1, 'A'), (0, 1, 'B'), (1, 2, 'C')])
+
+        def my_filter_function(edge):
+          if edge:
+            return edge == 'B'
+          return False  
+
+        indices = graph.filter_edges(my_filter_function)
+        print(indices)

src/digraph.rs

@@ -2843,6 +2843,84 @@ impl PyDiGraph {
         }
     }
 
+    /// Filters a graph's nodes by some criteria conditioned on a node's data payload and returns those nodes' indices.
+    ///
+    /// This function takes in a function as an argument. This filter function will be passed in a node's data payload and is
+    /// required to return a boolean value stating whether the node's data payload fits some criteria.
+    ///
+    /// For example::
+    ///     
+    ///     from rustworkx import PyDiGraph
+    ///
+    ///     graph = PyDiGraph()
+    ///     graph.add_nodes_from(list(range(5)))
+    ///
+    ///     def my_filter_function(node):
+    ///         return node > 2
+    ///
+    ///     indices = graph.filter_nodes(my_filter_function)
+    ///     assert indices == [3, 4]
+    ///
+    /// :param filter_function: Function with which to filter nodes
+    /// :returns: The node indices that match the filter
+    /// :rtype: NodeIndices
+    #[pyo3(text_signature = "(self, filter_function)")]
+    pub fn filter_nodes(&self, py: Python, filter_function: PyObject) -> PyResult<NodeIndices> {
+        let filter = |nindex: NodeIndex| -> PyResult<bool> {
+            let res = filter_function.call1(py, (&self.graph[nindex],))?;
+            res.extract(py)
+        };
+
+        let mut n = Vec::with_capacity(self.graph.node_count());
+        for node_index in self.graph.node_indices() {
+            if filter(node_index)? {
+                n.push(node_index.index())
+            };
+        }
+        Ok(NodeIndices { nodes: n })
+    }
+
+    /// Filters a graph's edges by some criteria conditioned on a edge's data payload and returns those edges' indices.
+    ///
+    /// This function takes in a function as an argument. This filter function will be passed in an edge's data payload and is
+    /// required to return a boolean value stating whether the edge's data payload fits some criteria.
+    ///
+    /// For example::
+    ///
+    ///     from rustworkx import PyGraph
+    ///     from rustworkx.generators import complete_graph
+    ///
+    ///     graph = PyGraph()
+    ///     graph.add_nodes_from(range(3))
+    ///     graph.add_edges_from([(0, 1, 'A'), (0, 1, 'B'), (1, 2, 'C')])
+    ///
+    ///     def my_filter_function(edge):
+    ///         if edge:
+    ///             return edge == 'B'
+    ///         return False  
+    ///        
+    ///     indices = graph.filter_edges(my_filter_function)
+    ///     assert indices == [1]
+    ///
+    /// :param filter_function: Function with which to filter edges
+    /// :returns: The edge indices that match the filter
+    /// :rtype: EdgeIndices
+    #[pyo3(text_signature = "(self, filter_function)")]
+    pub fn filter_edges(&self, py: Python, filter_function: PyObject) -> PyResult<EdgeIndices> {
+        let filter = |eindex: EdgeIndex| -> PyResult<bool> {
+            let res = filter_function.call1(py, (&self.graph[eindex],))?;
+            res.extract(py)
+        };
+
+        let mut e = Vec::with_capacity(self.graph.edge_count());
+        for edge_index in self.graph.edge_indices() {
+            if filter(edge_index)? {
+                e.push(edge_index.index())
+            };
+        }
+        Ok(EdgeIndices { edges: e })
+    }
+
     /// Return the number of nodes in the graph
     fn __len__(&self) -> PyResult<usize> {
         Ok(self.graph.node_count())

src/graph.rs

@@ -1848,6 +1848,84 @@ impl PyGraph {
         self.node_removed = false;
         self.attrs = py.None();
     }
+
+    /// Filters a graph's nodes by some criteria conditioned on a node's data payload and returns those nodes' indices.
+    ///
+    /// This function takes in a function as an argument. This filter function will be passed in a node's data payload and is
+    /// required to return a boolean value stating whether the node's data payload fits some criteria.
+    ///
+    /// For example::
+    ///     
+    ///     from rustworkx import PyGraph
+    ///
+    ///     graph = PyGraph()
+    ///     graph.add_nodes_from(list(range(5)))
+    ///
+    ///     def my_filter_function(node):
+    ///         return node > 2
+    ///
+    ///     indices = graph.filter_nodes(my_filter_function)
+    ///     assert indices == [3, 4]
+    ///
+    /// :param filter_function: Function with which to filter nodes
+    /// :returns: The node indices that match the filter
+    /// :rtype: NodeIndices
+    #[pyo3(text_signature = "(self, filter_function)")]
+    pub fn filter_nodes(&self, py: Python, filter_function: PyObject) -> PyResult<NodeIndices> {
+        let filter = |nindex: NodeIndex| -> PyResult<bool> {
+            let res = filter_function.call1(py, (&self.graph[nindex],))?;
+            res.extract(py)
+        };
+
+        let mut n = Vec::with_capacity(self.graph.node_count());
+        for node_index in self.graph.node_indices() {
+            if filter(node_index)? {
+                n.push(node_index.index())
+            };
+        }
+        Ok(NodeIndices { nodes: n })
+    }
+
+    /// Filters a graph's edges by some criteria conditioned on a edge's data payload and returns those edges' indices.
+    ///
+    /// This function takes in a function as an argument. This filter function will be passed in an edge's data payload and is
+    /// required to return a boolean value stating whether the edge's data payload fits some criteria.
+    ///
+    /// For example::
+    ///
+    ///     from rustworkx import PyGraph
+    ///     from rustworkx.generators import complete_graph
+    ///
+    ///     graph = PyGraph()
+    ///     graph.add_nodes_from(range(3))
+    ///     graph.add_edges_from([(0, 1, 'A'), (0, 1, 'B'), (1, 2, 'C')])
+    ///
+    ///     def my_filter_function(edge):
+    ///         if edge:
+    ///             return edge == 'B'
+    ///         return False  
+    ///        
+    ///     indices = graph.filter_edges(my_filter_function)
+    ///     assert indices == [1]
+    ///
+    /// :param filter_function: Function with which to filter edges
+    /// :returns: The edge indices that match the filter
+    /// :rtype: EdgeIndices
+    #[pyo3(text_signature = "(self, filter_function)")]
+    pub fn filter_edges(&self, py: Python, filter_function: PyObject) -> PyResult<EdgeIndices> {
+        let filter = |eindex: EdgeIndex| -> PyResult<bool> {
+            let res = filter_function.call1(py, (&self.graph[eindex],))?;
+            res.extract(py)
+        };
+
+        let mut e = Vec::with_capacity(self.graph.edge_count());
+        for edge_index in self.graph.edge_indices() {
+            if filter(edge_index)? {
+                e.push(edge_index.index())
+            };
+        }
+        Ok(EdgeIndices { edges: e })
+    }
 }
 
 fn weight_transform_callable(

tests/rustworkx_tests/digraph/test_filter.py

@@ -0,0 +1,81 @@
+# Licensed under the Apache License, Version 2.0 (the "License"); you may
+# not use this file except in compliance with the License. You may obtain
+# a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+# License for the specific language governing permissions and limitations
+# under the License.
+
+import unittest
+
+import rustworkx as rx
+
+
+class TestFilter(unittest.TestCase):
+    def test_filter_nodes(self):
+        def my_filter_function1(node):
+            return node == "cat"
+
+        def my_filter_function2(node):
+            return node == "lizard"
+
+        def my_filter_function3(node):
+            return node == "human"
+
+        graph = rx.PyDiGraph()
+        graph.add_node("cat")
+        graph.add_node("cat")
+        graph.add_node("dog")
+        graph.add_node("lizard")
+        graph.add_node("cat")
+        cat_indices = graph.filter_nodes(my_filter_function1)
+        lizard_indices = graph.filter_nodes(my_filter_function2)
+        human_indices = graph.filter_nodes(my_filter_function3)
+        self.assertEqual(list(cat_indices), [0, 1, 4])
+        self.assertEqual(list(lizard_indices), [3])
+        self.assertEqual(list(human_indices), [])
+
+    def test_filter_edges(self):
+        def my_filter_function1(edge):
+            return edge == "friends"
+
+        def my_filter_function2(edge):
+            return edge == "enemies"
+
+        def my_filter_function3(node):
+            return node == "frenemies"
+
+        graph = rx.PyDiGraph()
+        graph.add_node("cat")
+        graph.add_node("cat")
+        graph.add_node("dog")
+        graph.add_node("lizard")
+        graph.add_node("cat")
+        graph.add_edge(0, 2, "friends")
+        graph.add_edge(0, 1, "friends")
+        graph.add_edge(0, 3, "enemies")
+        friends_indices = graph.filter_edges(my_filter_function1)
+        enemies_indices = graph.filter_edges(my_filter_function2)
+        frenemies_indices = graph.filter_edges(my_filter_function3)
+        self.assertEqual(list(friends_indices), [0, 1])
+        self.assertEqual(list(enemies_indices), [2])
+        self.assertEqual(list(frenemies_indices), [])
+
+    def test_filter_errors(self):
+        def my_filter_function1(node):
+            raise TypeError("error!")
+
+        graph = rx.PyDiGraph()
+        graph.add_node("cat")
+        graph.add_node("cat")
+        graph.add_node("dog")
+        graph.add_edge(0, 1, "friends")
+        graph.add_edge(1, 2, "enemies")
+        with self.assertRaises(TypeError):
+            graph.filter_nodes(my_filter_function1)
+        with self.assertRaises(TypeError):
+            graph.filter_edges(my_filter_function1)

tests/rustworkx_tests/graph/test_filter.py

@@ -0,0 +1,81 @@
+# Licensed under the Apache License, Version 2.0 (the "License"); you may
+# not use this file except in compliance with the License. You may obtain
+# a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+# License for the specific language governing permissions and limitations
+# under the License.
+
+import unittest
+
+import rustworkx as rx
+
+
+class TestFilter(unittest.TestCase):
+    def test_filter_nodes(self):
+        def my_filter_function1(node):
+            return node == "cat"
+
+        def my_filter_function2(node):
+            return node == "lizard"
+
+        def my_filter_function3(node):
+            return node == "human"
+
+        graph = rx.PyGraph()
+        graph.add_node("cat")
+        graph.add_node("cat")
+        graph.add_node("dog")
+        graph.add_node("lizard")
+        graph.add_node("cat")
+        cat_indices = graph.filter_nodes(my_filter_function1)
+        lizard_indices = graph.filter_nodes(my_filter_function2)
+        human_indices = graph.filter_nodes(my_filter_function3)
+        self.assertEqual(list(cat_indices), [0, 1, 4])
+        self.assertEqual(list(lizard_indices), [3])
+        self.assertEqual(list(human_indices), [])
+
+    def test_filter_edges(self):
+        def my_filter_function1(edge):
+            return edge == "friends"
+
+        def my_filter_function2(edge):
+            return edge == "enemies"
+
+        def my_filter_function3(node):
+            return node == "frenemies"
+
+        graph = rx.PyGraph()
+        graph.add_node("cat")
+        graph.add_node("cat")
+        graph.add_node("dog")
+        graph.add_node("lizard")
+        graph.add_node("cat")
+        graph.add_edge(0, 2, "friends")
+        graph.add_edge(0, 1, "friends")
+        graph.add_edge(0, 3, "enemies")
+        friends_indices = graph.filter_edges(my_filter_function1)
+        enemies_indices = graph.filter_edges(my_filter_function2)
+        frenemies_indices = graph.filter_edges(my_filter_function3)
+        self.assertEqual(list(friends_indices), [0, 1])
+        self.assertEqual(list(enemies_indices), [2])
+        self.assertEqual(list(frenemies_indices), [])
+
+    def test_filter_errors(self):
+        def my_filter_function1(node):
+            raise TypeError("error!")
+
+        graph = rx.PyGraph()
+        graph.add_node("cat")
+        graph.add_node("cat")
+        graph.add_node("dog")
+        graph.add_edge(0, 1, "friends")
+        graph.add_edge(1, 2, "enemies")
+        with self.assertRaises(TypeError):
+            graph.filter_nodes(my_filter_function1)
+        with self.assertRaises(TypeError):
+            graph.filter_edges(my_filter_function1)