Add the equivalent code from the previous post

ajavageek · Oct 18, 2023 · 7ad74a5 · 7ad74a5
1 parent 1da2e41
commit 7ad74a5
Show file tree

Hide file tree

Showing 3 changed files with 59 additions and 18 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -9,4 +9,5 @@ name = "rust_over_pyo3"
 crate-type = ["cdylib"]
 
 [dependencies]
-pyo3 = "0.19.0"
+pyo3 = { version = "0.20.0" , features = ["num-complex"]}
+num-complex = "0.4.4"
diff --git a/src/lib.rs b/src/lib.rs
@@ -1,3 +1,5 @@
+use num_complex::Complex;
+use pyo3::exceptions::PyValueError;
 use pyo3::prelude::*;
 
 /// Formats the sum of two numbers as string.
@@ -10,5 +12,16 @@ fn sum_as_string(a: usize, b: usize) -> PyResult<String> {
 #[pymodule]
 fn rust_over_pyo3(_py: Python, m: &PyModule) -> PyResult<()> {
     m.add_function(wrap_pyfunction!(sum_as_string, m)?)?;
+    m.add_function(wrap_pyfunction!(compute, m)?)?;
     Ok(())
-}
+}
+
+#[pyfunction]
+fn compute(command: &str, a: Complex<f64>, b: Complex<f64>) -> PyResult<Complex<f64>> {
+    match command {
+        "add" => Ok(a + b),
+        "sub" => Ok(a - b),
+        "mul" => Ok(a * b),
+        _ => Err(PyValueError::new_err("Unknown command")),
+    }
+}