Remove repeated weight_callable code (#462)

Profiting from #459 that made weight_callable generic to remove some code duplication
Qiskit · Oct 15, 2021 · 9dff1bb · 9dff1bb
1 parent ba02c6b
commit 9dff1bb
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 27 deletions.
diff --git a/src/digraph.rs b/src/digraph.rs
@@ -50,8 +50,8 @@ use super::iterators::{
     EdgeIndexMap, EdgeIndices, EdgeList, NodeIndices, NodeMap, WeightedEdgeList,
 };
 use super::{
-    find_node_by_weight, DAGHasCycle, DAGWouldCycle, IsNan, NoEdgeBetweenNodes,
-    NoSuitableNeighbors, NodesRemoved,
+    find_node_by_weight, weight_callable, DAGHasCycle, DAGWouldCycle, IsNan,
+    NoEdgeBetweenNodes, NoSuitableNeighbors, NodesRemoved,
 };
 
 use super::dag_algo::is_directed_acyclic_graph;
@@ -1903,17 +1903,6 @@ impl PyDiGraph {
             None => " ".to_string(),
         };
 
-        let weight_callable = |value: &PyObject,
-                               weight_fn: &Option<PyObject>|
-         -> PyResult<Option<String>> {
-            match weight_fn {
-                Some(weight_fn) => {
-                    let res = weight_fn.call1(py, (value,))?;
-                    Ok(Some(res.extract(py)?))
-                }
-                None => Ok(None),
-            }
-        };
         for edge in self.graph.edge_references() {
             buf_writer.write_all(
                 format!(
@@ -1924,7 +1913,12 @@ impl PyDiGraph {
                 )
                 .as_bytes(),
             )?;
-            match weight_callable(edge.weight(), &weight_fn)? {
+            match weight_callable(
+                py,
+                &weight_fn,
+                edge.weight(),
+                None as Option<String>,
+            )? {
                 Some(weight) => buf_writer
                     .write_all(format!("{}{}\n", delim, weight).as_bytes()),
                 None => buf_writer.write_all(b"\n"),

diff --git a/src/graph.rs b/src/graph.rs
@@ -37,7 +37,10 @@ use super::dot_utils::build_dot;
 use super::iterators::{
     EdgeIndexMap, EdgeIndices, EdgeList, NodeIndices, WeightedEdgeList,
 };
-use super::{find_node_by_weight, IsNan, NoEdgeBetweenNodes, NodesRemoved};
+use super::{
+    find_node_by_weight, weight_callable, IsNan, NoEdgeBetweenNodes,
+    NodesRemoved,
+};
 
 use petgraph::algo;
 use petgraph::graph::{EdgeIndex, NodeIndex};
@@ -1259,17 +1262,6 @@ impl PyGraph {
             None => " ".to_string(),
         };
 
-        let weight_callable = |value: &PyObject,
-                               weight_fn: &Option<PyObject>|
-         -> PyResult<Option<String>> {
-            match weight_fn {
-                Some(weight_fn) => {
-                    let res = weight_fn.call1(py, (value,))?;
-                    Ok(Some(res.extract(py)?))
-                }
-                None => Ok(None),
-            }
-        };
         for edge in self.graph.edge_references() {
             buf_writer.write_all(
                 format!(
@@ -1280,7 +1272,12 @@ impl PyGraph {
                 )
                 .as_bytes(),
             )?;
-            match weight_callable(edge.weight(), &weight_fn)? {
+            match weight_callable(
+                py,
+                &weight_fn,
+                edge.weight(),
+                None as Option<String>,
+            )? {
                 Some(weight) => buf_writer
                     .write_all(format!("{}{}\n", delim, weight).as_bytes()),
                 None => buf_writer.write_all(b"\n"),