pyo3/types/
complex.rs

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