Fix: Return an Iterator instance from layers in rustworkx-core.

Qiskit · May 17, 2024 · 33e6bbc · 33e6bbc
1 parent 7da636f
commit 33e6bbc
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 11 deletions.
diff --git a/rustworkx-core/src/layers.rs b/rustworkx-core/src/layers.rs
@@ -15,6 +15,7 @@ use petgraph::visit::{
     EdgeRef, GraphBase, IntoEdgesDirected, IntoNeighborsDirected, IntoNodeIdentifiers,
     NodeIndexable,
 };
+use std::{error::Error, fmt, hash::Hash, vec::IntoIter};
 
 /// Error returned by Layers function when an index is not part of the graph.
 #[derive(Debug, PartialEq, Eq)]
@@ -34,7 +35,7 @@ impl fmt::Display for LayersInvalidIndex {
     }
 }
 
-use std::{error::Error, fmt, hash::Hash};
+pub type LayersIterator = IntoIter<Vec<usize>>;
 
 /// Return a list of graph layers
 ///
@@ -60,11 +61,11 @@ use std::{error::Error, fmt, hash::Hash};
 /// ];
 ///
 /// let graph = DiGraph::<u32, u32>::from_edges(&edge_list);
-/// let layers = layers(&graph, vec![0,1]).unwrap();
+/// let layers: Vec<Vec<usize>> = layers(&graph, vec![0,1]).unwrap().collect();
 /// let expected_layers = vec![vec![0,1], vec![1,2], vec![2,3], vec![3,4], vec![4]];
 /// assert_eq!(layers, expected_layers)
 /// ```
-pub fn layers<G>(graph: G, first_layer: Vec<usize>) -> Result<Vec<Vec<usize>>, LayersInvalidIndex>
+pub fn layers<G>(graph: G, first_layer: Vec<usize>) -> Result<LayersIterator, LayersInvalidIndex>
 where
     G: NodeIndexable // Used in from_index and to_index.
         + IntoNodeIdentifiers // Used for .node_identifiers
@@ -132,5 +133,5 @@ where
         cur_layer = next_layer;
         next_layer = Vec::new();
     }
-    Ok(output_indices)
+    Ok(output_indices.into_iter())
 }
diff --git a/src/dag_algo/mod.rs b/src/dag_algo/mod.rs
@@ -256,18 +256,18 @@ pub fn layers(
     match result {
         Ok(result) => {
             if index_output {
-                Ok(result.to_object(py))
+                Ok(PyList::new_bound(py, result).into())
             } else {
-                Ok(result
-                    .iter()
-                    .map(|x| {
+                Ok(PyList::new_bound(
+                    py,
+                    result.map(|x| {
                         x.iter()
                             .map(|index| dag.graph.node_weight(dag.graph.from_index(*index)))
                             .collect::<Vec<Option<&PyObject>>>()
                             .to_object(py)
-                    })
-                    .collect::<Vec<PyObject>>()
-                    .to_object(py))
+                    }),
+                )
+                .into())
             }
         }
         Err(e) => Err(InvalidNode::new_err(e.0)),