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    #checkfunction=ffi::PyComplex_Check
31);
32
33impl PyComplex {
34    /// Creates a new `PyComplex` from the given real and imaginary values.
35    pub fn from_doubles(py: Python<'_>, real: c_double, imag: c_double) -> Bound<'_, PyComplex> {
36        use crate::ffi_ptr_ext::FfiPtrExt;
37        unsafe {
38            ffi::PyComplex_FromDoubles(real, imag)
39                .assume_owned(py)
40                .cast_into_unchecked()
41        }
42    }
43}
44
45#[cfg(not(any(Py_LIMITED_API, PyPy, GraalPy)))]
46mod not_limited_impls {
47    use crate::Borrowed;
48
49    use super::*;
50    use std::ops::{Add, Div, Mul, Neg, Sub};
51
52    macro_rules! bin_ops {
53        ($trait:ident, $fn:ident, $op:tt) => {
54            impl<'py> $trait for Borrowed<'_, 'py, PyComplex> {
55                type Output = Bound<'py, PyComplex>;
56                fn $fn(self, other: Self) -> Self::Output {
57                    PyAnyMethods::$fn(self.as_any(), other)
58                    .cast_into().expect(
59                        concat!("Complex method ",
60                            stringify!($fn),
61                            " failed.")
62                        )
63                }
64            }
65
66            impl<'py> $trait for &Bound<'py, PyComplex> {
67                type Output = Bound<'py, PyComplex>;
68                fn $fn(self, other: &Bound<'py, PyComplex>) -> Bound<'py, PyComplex> {
69                    self.as_borrowed() $op other.as_borrowed()
70                }
71            }
72
73            impl<'py> $trait<Bound<'py, PyComplex>> for &Bound<'py, PyComplex> {
74                type Output = Bound<'py, PyComplex>;
75                fn $fn(self, other: Bound<'py, PyComplex>) -> Bound<'py, PyComplex> {
76                    self.as_borrowed() $op other.as_borrowed()
77                }
78            }
79
80            impl<'py> $trait for Bound<'py, PyComplex> {
81                type Output = Bound<'py, PyComplex>;
82                fn $fn(self, other: Bound<'py, PyComplex>) -> Bound<'py, PyComplex> {
83                    self.as_borrowed() $op other.as_borrowed()
84                }
85            }
86
87            impl<'py> $trait<&Self> for Bound<'py, PyComplex> {
88                type Output = Bound<'py, PyComplex>;
89                fn $fn(self, other: &Bound<'py, PyComplex>) -> Bound<'py, PyComplex> {
90                    self.as_borrowed() $op other.as_borrowed()
91                }
92            }
93        };
94    }
95
96    bin_ops!(Add, add, +);
97    bin_ops!(Sub, sub, -);
98    bin_ops!(Mul, mul, *);
99    bin_ops!(Div, div, /);
100
101    impl<'py> Neg for Borrowed<'_, 'py, PyComplex> {
102        type Output = Bound<'py, PyComplex>;
103        fn neg(self) -> Self::Output {
104            PyAnyMethods::neg(self.as_any())
105                .cast_into()
106                .expect("Complex method __neg__ failed.")
107        }
108    }
109
110    impl<'py> Neg for &Bound<'py, PyComplex> {
111        type Output = Bound<'py, PyComplex>;
112        fn neg(self) -> Bound<'py, PyComplex> {
113            -self.as_borrowed()
114        }
115    }
116
117    impl<'py> Neg for Bound<'py, PyComplex> {
118        type Output = Bound<'py, PyComplex>;
119        fn neg(self) -> Bound<'py, PyComplex> {
120            -self.as_borrowed()
121        }
122    }
123
124    #[cfg(test)]
125    mod tests {
126        use super::PyComplex;
127        use crate::{types::complex::PyComplexMethods, Python};
128        use assert_approx_eq::assert_approx_eq;
129
130        #[test]
131        fn test_add() {
132            Python::attach(|py| {
133                let l = PyComplex::from_doubles(py, 3.0, 1.2);
134                let r = PyComplex::from_doubles(py, 1.0, 2.6);
135                let res = l + r;
136                assert_approx_eq!(res.real(), 4.0);
137                assert_approx_eq!(res.imag(), 3.8);
138            });
139        }
140
141        #[test]
142        fn test_sub() {
143            Python::attach(|py| {
144                let l = PyComplex::from_doubles(py, 3.0, 1.2);
145                let r = PyComplex::from_doubles(py, 1.0, 2.6);
146                let res = l - r;
147                assert_approx_eq!(res.real(), 2.0);
148                assert_approx_eq!(res.imag(), -1.4);
149            });
150        }
151
152        #[test]
153        fn test_mul() {
154            Python::attach(|py| {
155                let l = PyComplex::from_doubles(py, 3.0, 1.2);
156                let r = PyComplex::from_doubles(py, 1.0, 2.6);
157                let res = l * r;
158                assert_approx_eq!(res.real(), -0.12);
159                assert_approx_eq!(res.imag(), 9.0);
160            });
161        }
162
163        #[test]
164        fn test_div() {
165            Python::attach(|py| {
166                let l = PyComplex::from_doubles(py, 3.0, 1.2);
167                let r = PyComplex::from_doubles(py, 1.0, 2.6);
168                let res = l / r;
169                assert_approx_eq!(res.real(), 0.788_659_793_814_432_9);
170                assert_approx_eq!(res.imag(), -0.850_515_463_917_525_7);
171            });
172        }
173
174        #[test]
175        fn test_neg() {
176            Python::attach(|py| {
177                let val = PyComplex::from_doubles(py, 3.0, 1.2);
178                let res = -val;
179                assert_approx_eq!(res.real(), -3.0);
180                assert_approx_eq!(res.imag(), -1.2);
181            });
182        }
183
184        #[test]
185        fn test_abs() {
186            Python::attach(|py| {
187                let val = PyComplex::from_doubles(py, 3.0, 1.2);
188                assert_approx_eq!(val.abs(), 3.231_098_884_280_702_2);
189            });
190        }
191
192        #[test]
193        fn test_pow() {
194            Python::attach(|py| {
195                let l = PyComplex::from_doubles(py, 3.0, 1.2);
196                let r = PyComplex::from_doubles(py, 1.2, 2.6);
197                let val = l.pow(&r);
198                assert_approx_eq!(val.real(), -1.419_309_997_016_603_7);
199                assert_approx_eq!(val.imag(), -0.541_297_466_033_544_6);
200            });
201        }
202    }
203}
204
205/// Implementation of functionality for [`PyComplex`].
206///
207/// These methods are defined for the `Bound<'py, PyComplex>` smart pointer, so to use method call
208/// syntax these methods are separated into a trait, because stable Rust does not yet support
209/// `arbitrary_self_types`.
210#[doc(alias = "PyComplex")]
211pub trait PyComplexMethods<'py>: crate::sealed::Sealed {
212    /// Returns the real part of the complex number.
213    fn real(&self) -> c_double;
214    /// Returns the imaginary part of the complex number.
215    fn imag(&self) -> c_double;
216    /// Returns `|self|`.
217    #[cfg(not(any(Py_LIMITED_API, PyPy, GraalPy)))]
218    fn abs(&self) -> c_double;
219    /// Returns `self` raised to the power of `other`.
220    #[cfg(not(any(Py_LIMITED_API, PyPy, GraalPy)))]
221    fn pow(&self, other: &Bound<'py, PyComplex>) -> Bound<'py, PyComplex>;
222}
223
224impl<'py> PyComplexMethods<'py> for Bound<'py, PyComplex> {
225    fn real(&self) -> c_double {
226        unsafe { ffi::PyComplex_RealAsDouble(self.as_ptr()) }
227    }
228
229    fn imag(&self) -> c_double {
230        unsafe { ffi::PyComplex_ImagAsDouble(self.as_ptr()) }
231    }
232
233    #[cfg(not(any(Py_LIMITED_API, PyPy, GraalPy)))]
234    fn abs(&self) -> c_double {
235        PyAnyMethods::abs(self.as_any())
236            .cast_into()
237            .expect("Complex method __abs__ failed.")
238            .extract()
239            .expect("Failed to extract to c double.")
240    }
241
242    #[cfg(not(any(Py_LIMITED_API, PyPy, GraalPy)))]
243    fn pow(&self, other: &Bound<'py, PyComplex>) -> Bound<'py, PyComplex> {
244        Python::attach(|py| {
245            PyAnyMethods::pow(self.as_any(), other, py.None())
246                .cast_into()
247                .expect("Complex method __pow__ failed.")
248        })
249    }
250}
251
252#[cfg(test)]
253mod tests {
254    use super::PyComplex;
255    use crate::{types::complex::PyComplexMethods, Python};
256    use assert_approx_eq::assert_approx_eq;
257
258    #[test]
259    fn test_from_double() {
260        use assert_approx_eq::assert_approx_eq;
261
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
