pyo3/types/
complex.rs

1#[cfg(not(any(Py_LIMITED_API, PyPy, GraalPy)))]
2use crate::py_result_ext::PyResultExt;
3#[cfg(not(any(Py_LIMITED_API, PyPy, GraalPy)))]
4use crate::types::any::PyAnyMethods;
5use crate::{ffi, Bound, PyAny, Python};
6use std::ffi::c_double;
7
8/// Represents a Python [`complex`](https://docs.python.org/3/library/functions.html#complex) object.
9///
10/// Values of this type are accessed via PyO3's smart pointers, e.g. as
11/// [`Py<PyComplex>`][crate::Py] or [`Bound<'py, PyComplex>`][Bound].
12///
13/// For APIs available on `complex` objects, see the [`PyComplexMethods`] trait which is implemented for
14/// [`Bound<'py, PyComplex>`][Bound].
15///
16/// Note that `PyComplex` supports only basic operations. For advanced operations
17/// consider using [num-complex](https://docs.rs/num-complex)'s [`Complex`] type instead.
18/// This optional dependency can be activated with the `num-complex` feature flag.
19///
20/// [`Complex`]: https://docs.rs/num-complex/latest/num_complex/struct.Complex.html
21#[repr(transparent)]
22pub struct PyComplex(PyAny);
23
24pyobject_subclassable_native_type!(PyComplex, ffi::PyComplexObject);
25
26pyobject_native_type!(
27    PyComplex,
28    ffi::PyComplexObject,
29    pyobject_native_static_type_object!(ffi::PyComplex_Type),
30    "builtins",
31    "complex",
32    #checkfunction=ffi::PyComplex_Check
33);
34
35impl PyComplex {
36    /// Creates a new `PyComplex` from the given real and imaginary values.
37    pub fn from_doubles(py: Python<'_>, real: c_double, imag: c_double) -> Bound<'_, PyComplex> {
38        use crate::ffi_ptr_ext::FfiPtrExt;
39        unsafe {
40            ffi::PyComplex_FromDoubles(real, imag)
41                .assume_owned(py)
42                .cast_into_unchecked()
43        }
44    }
45}
46
47#[cfg(not(any(Py_LIMITED_API, PyPy, GraalPy)))]
48mod not_limited_impls {
49    use crate::Borrowed;
50
51    use super::*;
52    use std::ops::{Add, Div, Mul, Neg, Sub};
53
54    macro_rules! bin_ops {
55        ($trait:ident, $fn:ident, $op:tt) => {
56            impl<'py> $trait for Borrowed<'_, 'py, PyComplex> {
57                type Output = Bound<'py, PyComplex>;
58                fn $fn(self, other: Self) -> Self::Output {
59                    PyAnyMethods::$fn(self.as_any(), other)
60                    .cast_into().expect(
61                        concat!("Complex method ",
62                            stringify!($fn),
63                            " failed.")
64                        )
65                }
66            }
67
68            impl<'py> $trait for &Bound<'py, PyComplex> {
69                type Output = Bound<'py, PyComplex>;
70                fn $fn(self, other: &Bound<'py, PyComplex>) -> Bound<'py, PyComplex> {
71                    self.as_borrowed() $op other.as_borrowed()
72                }
73            }
74
75            impl<'py> $trait<Bound<'py, PyComplex>> for &Bound<'py, PyComplex> {
76                type Output = Bound<'py, PyComplex>;
77                fn $fn(self, other: Bound<'py, PyComplex>) -> Bound<'py, PyComplex> {
78                    self.as_borrowed() $op other.as_borrowed()
79                }
80            }
81
82            impl<'py> $trait for Bound<'py, PyComplex> {
83                type Output = Bound<'py, PyComplex>;
84                fn $fn(self, other: Bound<'py, PyComplex>) -> Bound<'py, PyComplex> {
85                    self.as_borrowed() $op other.as_borrowed()
86                }
87            }
88
89            impl<'py> $trait<&Self> for Bound<'py, PyComplex> {
90                type Output = Bound<'py, PyComplex>;
91                fn $fn(self, other: &Bound<'py, PyComplex>) -> Bound<'py, PyComplex> {
92                    self.as_borrowed() $op other.as_borrowed()
93                }
94            }
95        };
96    }
97
98    bin_ops!(Add, add, +);
99    bin_ops!(Sub, sub, -);
100    bin_ops!(Mul, mul, *);
101    bin_ops!(Div, div, /);
102
103    impl<'py> Neg for Borrowed<'_, 'py, PyComplex> {
104        type Output = Bound<'py, PyComplex>;
105        fn neg(self) -> Self::Output {
106            PyAnyMethods::neg(self.as_any())
107                .cast_into()
108                .expect("Complex method __neg__ failed.")
109        }
110    }
111
112    impl<'py> Neg for &Bound<'py, PyComplex> {
113        type Output = Bound<'py, PyComplex>;
114        fn neg(self) -> Bound<'py, PyComplex> {
115            -self.as_borrowed()
116        }
117    }
118
119    impl<'py> Neg for Bound<'py, PyComplex> {
120        type Output = Bound<'py, PyComplex>;
121        fn neg(self) -> Bound<'py, PyComplex> {
122            -self.as_borrowed()
123        }
124    }
125
126    #[cfg(test)]
127    mod tests {
128        use super::PyComplex;
129        use crate::{types::complex::PyComplexMethods, Python};
130        use assert_approx_eq::assert_approx_eq;
131
132        #[test]
133        fn test_add() {
134            Python::attach(|py| {
135                let l = PyComplex::from_doubles(py, 3.0, 1.2);
136                let r = PyComplex::from_doubles(py, 1.0, 2.6);
137                let res = l + r;
138                assert_approx_eq!(res.real(), 4.0);
139                assert_approx_eq!(res.imag(), 3.8);
140            });
141        }
142
143        #[test]
144        fn test_sub() {
145            Python::attach(|py| {
146                let l = PyComplex::from_doubles(py, 3.0, 1.2);
147                let r = PyComplex::from_doubles(py, 1.0, 2.6);
148                let res = l - r;
149                assert_approx_eq!(res.real(), 2.0);
150                assert_approx_eq!(res.imag(), -1.4);
151            });
152        }
153
154        #[test]
155        fn test_mul() {
156            Python::attach(|py| {
157                let l = PyComplex::from_doubles(py, 3.0, 1.2);
158                let r = PyComplex::from_doubles(py, 1.0, 2.6);
159                let res = l * r;
160                assert_approx_eq!(res.real(), -0.12);
161                assert_approx_eq!(res.imag(), 9.0);
162            });
163        }
164
165        #[test]
166        fn test_div() {
167            Python::attach(|py| {
168                let l = PyComplex::from_doubles(py, 3.0, 1.2);
169                let r = PyComplex::from_doubles(py, 1.0, 2.6);
170                let res = l / r;
171                assert_approx_eq!(res.real(), 0.788_659_793_814_432_9);
172                assert_approx_eq!(res.imag(), -0.850_515_463_917_525_7);
173            });
174        }
175
176        #[test]
177        fn test_neg() {
178            Python::attach(|py| {
179                let val = PyComplex::from_doubles(py, 3.0, 1.2);
180                let res = -val;
181                assert_approx_eq!(res.real(), -3.0);
182                assert_approx_eq!(res.imag(), -1.2);
183            });
184        }
185
186        #[test]
187        fn test_abs() {
188            Python::attach(|py| {
189                let val = PyComplex::from_doubles(py, 3.0, 1.2);
190                assert_approx_eq!(val.abs(), 3.231_098_884_280_702_2);
191            });
192        }
193
194        #[test]
195        fn test_pow() {
196            Python::attach(|py| {
197                let l = PyComplex::from_doubles(py, 3.0, 1.2);
198                let r = PyComplex::from_doubles(py, 1.2, 2.6);
199                let val = l.pow(&r);
200                assert_approx_eq!(val.real(), -1.419_309_997_016_603_7);
201                assert_approx_eq!(val.imag(), -0.541_297_466_033_544_6);
202            });
203        }
204    }
205}
206
207/// Implementation of functionality for [`PyComplex`].
208///
209/// These methods are defined for the `Bound<'py, PyComplex>` smart pointer, so to use method call
210/// syntax these methods are separated into a trait, because stable Rust does not yet support
211/// `arbitrary_self_types`.
212#[doc(alias = "PyComplex")]
213pub trait PyComplexMethods<'py>: crate::sealed::Sealed {
214    /// Returns the real part of the complex number.
215    fn real(&self) -> c_double;
216    /// Returns the imaginary part of the complex number.
217    fn imag(&self) -> c_double;
218    /// Returns `|self|`.
219    #[cfg(not(any(Py_LIMITED_API, PyPy, GraalPy)))]
220    fn abs(&self) -> c_double;
221    /// Returns `self` raised to the power of `other`.
222    #[cfg(not(any(Py_LIMITED_API, PyPy, GraalPy)))]
223    fn pow(&self, other: &Bound<'py, PyComplex>) -> Bound<'py, PyComplex>;
224}
225
226impl<'py> PyComplexMethods<'py> for Bound<'py, PyComplex> {
227    fn real(&self) -> c_double {
228        unsafe { ffi::PyComplex_RealAsDouble(self.as_ptr()) }
229    }
230
231    fn imag(&self) -> c_double {
232        unsafe { ffi::PyComplex_ImagAsDouble(self.as_ptr()) }
233    }
234
235    #[cfg(not(any(Py_LIMITED_API, PyPy, GraalPy)))]
236    fn abs(&self) -> c_double {
237        PyAnyMethods::abs(self.as_any())
238            .cast_into()
239            .expect("Complex method __abs__ failed.")
240            .extract()
241            .expect("Failed to extract to c double.")
242    }
243
244    #[cfg(not(any(Py_LIMITED_API, PyPy, GraalPy)))]
245    fn pow(&self, other: &Bound<'py, PyComplex>) -> Bound<'py, PyComplex> {
246        Python::attach(|py| {
247            PyAnyMethods::pow(self.as_any(), other, py.None())
248                .cast_into()
249                .expect("Complex method __pow__ failed.")
250        })
251    }
252}
253
254#[cfg(test)]
255mod tests {
256    use super::PyComplex;
257    use crate::{types::complex::PyComplexMethods, Python};
258    use assert_approx_eq::assert_approx_eq;
259
260    #[test]
261    fn test_from_double() {
262        Python::attach(|py| {
263            let complex = PyComplex::from_doubles(py, 3.0, 1.2);
264            assert_approx_eq!(complex.real(), 3.0);
265            assert_approx_eq!(complex.imag(), 1.2);
266        });
267    }
268}
pyo3/types/complex.rs

pyo3/types/
complex.rs
