pyo3/types/
complex.rs

1#[cfg(not(any(Py_LIMITED_API, PyPy, GraalPy)))]
2use crate::py_result_ext::PyResultExt;
3#[cfg(feature = "gil-refs")]
4use crate::PyNativeType;
5use crate::{ffi, types::any::PyAnyMethods, Bound, PyAny, Python};
6use std::os::raw::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_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_bound(
34        py: Python<'_>,
35        real: c_double,
36        imag: c_double,
37    ) -> Bound<'_, PyComplex> {
38        use crate::ffi_ptr_ext::FfiPtrExt;
39        unsafe {
40            ffi::PyComplex_FromDoubles(real, imag)
41                .assume_owned(py)
42                .downcast_into_unchecked()
43        }
44    }
45}
46
47#[cfg(feature = "gil-refs")]
48impl PyComplex {
49    /// Deprecated form of [`PyComplex::from_doubles_bound`]
50    #[deprecated(
51        since = "0.21.0",
52        note = "`PyComplex::from_doubles` will be replaced by `PyComplex::from_doubles_bound` in a future PyO3 version"
53    )]
54    pub fn from_doubles(py: Python<'_>, real: c_double, imag: c_double) -> &PyComplex {
55        Self::from_doubles_bound(py, real, imag).into_gil_ref()
56    }
57
58    /// Returns the real part of the complex number.
59    pub fn real(&self) -> c_double {
60        self.as_borrowed().real()
61    }
62    /// Returns the imaginary part of the complex number.
63    pub fn imag(&self) -> c_double {
64        self.as_borrowed().imag()
65    }
66}
67
68#[cfg(not(any(Py_LIMITED_API, PyPy, GraalPy)))]
69mod not_limited_impls {
70    use crate::Borrowed;
71
72    use super::*;
73    use std::ops::{Add, Div, Mul, Neg, Sub};
74
75    #[cfg(feature = "gil-refs")]
76    impl PyComplex {
77        /// Returns `|self|`.
78        pub fn abs(&self) -> c_double {
79            self.as_borrowed().abs()
80        }
81        /// Returns `self` raised to the power of `other`.
82        pub fn pow<'py>(&'py self, other: &'py PyComplex) -> &'py PyComplex {
83            self.as_borrowed().pow(&other.as_borrowed()).into_gil_ref()
84        }
85    }
86
87    macro_rules! bin_ops {
88        ($trait:ident, $fn:ident, $op:tt) => {
89            impl<'py> $trait for Borrowed<'_, 'py, PyComplex> {
90                type Output = Bound<'py, PyComplex>;
91                fn $fn(self, other: Self) -> Self::Output {
92                    PyAnyMethods::$fn(self.as_any(), other)
93                    .downcast_into().expect(
94                        concat!("Complex method ",
95                            stringify!($fn),
96                            " failed.")
97                        )
98                }
99            }
100
101            #[cfg(feature = "gil-refs")]
102            impl<'py> $trait for &'py PyComplex {
103                type Output = &'py PyComplex;
104                fn $fn(self, other: &'py PyComplex) -> &'py PyComplex {
105                    (self.as_borrowed() $op other.as_borrowed()).into_gil_ref()
106                }
107            }
108
109            impl<'py> $trait for &Bound<'py, PyComplex> {
110                type Output = Bound<'py, PyComplex>;
111                fn $fn(self, other: &Bound<'py, PyComplex>) -> Bound<'py, PyComplex> {
112                    self.as_borrowed() $op other.as_borrowed()
113                }
114            }
115
116            impl<'py> $trait<Bound<'py, PyComplex>> for &Bound<'py, PyComplex> {
117                type Output = Bound<'py, PyComplex>;
118                fn $fn(self, other: Bound<'py, PyComplex>) -> Bound<'py, PyComplex> {
119                    self.as_borrowed() $op other.as_borrowed()
120                }
121            }
122
123            impl<'py> $trait for Bound<'py, PyComplex> {
124                type Output = Bound<'py, PyComplex>;
125                fn $fn(self, other: Bound<'py, PyComplex>) -> Bound<'py, PyComplex> {
126                    self.as_borrowed() $op other.as_borrowed()
127                }
128            }
129
130            impl<'py> $trait<&Self> for Bound<'py, PyComplex> {
131                type Output = Bound<'py, PyComplex>;
132                fn $fn(self, other: &Bound<'py, PyComplex>) -> Bound<'py, PyComplex> {
133                    self.as_borrowed() $op other.as_borrowed()
134                }
135            }
136        };
137    }
138
139    bin_ops!(Add, add, +);
140    bin_ops!(Sub, sub, -);
141    bin_ops!(Mul, mul, *);
142    bin_ops!(Div, div, /);
143
144    #[cfg(feature = "gil-refs")]
145    impl<'py> Neg for &'py PyComplex {
146        type Output = &'py PyComplex;
147        fn neg(self) -> &'py PyComplex {
148            (-self.as_borrowed()).into_gil_ref()
149        }
150    }
151
152    impl<'py> Neg for Borrowed<'_, 'py, PyComplex> {
153        type Output = Bound<'py, PyComplex>;
154        fn neg(self) -> Self::Output {
155            PyAnyMethods::neg(self.as_any())
156                .downcast_into()
157                .expect("Complex method __neg__ failed.")
158        }
159    }
160
161    impl<'py> Neg for &Bound<'py, PyComplex> {
162        type Output = Bound<'py, PyComplex>;
163        fn neg(self) -> Bound<'py, PyComplex> {
164            -self.as_borrowed()
165        }
166    }
167
168    impl<'py> Neg for Bound<'py, PyComplex> {
169        type Output = Bound<'py, PyComplex>;
170        fn neg(self) -> Bound<'py, PyComplex> {
171            -self.as_borrowed()
172        }
173    }
174
175    #[cfg(test)]
176    mod tests {
177        use super::PyComplex;
178        use crate::{types::complex::PyComplexMethods, Python};
179        use assert_approx_eq::assert_approx_eq;
180
181        #[test]
182        fn test_add() {
183            Python::with_gil(|py| {
184                let l = PyComplex::from_doubles_bound(py, 3.0, 1.2);
185                let r = PyComplex::from_doubles_bound(py, 1.0, 2.6);
186                let res = l + r;
187                assert_approx_eq!(res.real(), 4.0);
188                assert_approx_eq!(res.imag(), 3.8);
189            });
190        }
191
192        #[test]
193        fn test_sub() {
194            Python::with_gil(|py| {
195                let l = PyComplex::from_doubles_bound(py, 3.0, 1.2);
196                let r = PyComplex::from_doubles_bound(py, 1.0, 2.6);
197                let res = l - r;
198                assert_approx_eq!(res.real(), 2.0);
199                assert_approx_eq!(res.imag(), -1.4);
200            });
201        }
202
203        #[test]
204        fn test_mul() {
205            Python::with_gil(|py| {
206                let l = PyComplex::from_doubles_bound(py, 3.0, 1.2);
207                let r = PyComplex::from_doubles_bound(py, 1.0, 2.6);
208                let res = l * r;
209                assert_approx_eq!(res.real(), -0.12);
210                assert_approx_eq!(res.imag(), 9.0);
211            });
212        }
213
214        #[test]
215        fn test_div() {
216            Python::with_gil(|py| {
217                let l = PyComplex::from_doubles_bound(py, 3.0, 1.2);
218                let r = PyComplex::from_doubles_bound(py, 1.0, 2.6);
219                let res = l / r;
220                assert_approx_eq!(res.real(), 0.788_659_793_814_432_9);
221                assert_approx_eq!(res.imag(), -0.850_515_463_917_525_7);
222            });
223        }
224
225        #[test]
226        fn test_neg() {
227            Python::with_gil(|py| {
228                let val = PyComplex::from_doubles_bound(py, 3.0, 1.2);
229                let res = -val;
230                assert_approx_eq!(res.real(), -3.0);
231                assert_approx_eq!(res.imag(), -1.2);
232            });
233        }
234
235        #[test]
236        fn test_abs() {
237            Python::with_gil(|py| {
238                let val = PyComplex::from_doubles_bound(py, 3.0, 1.2);
239                assert_approx_eq!(val.abs(), 3.231_098_884_280_702_2);
240            });
241        }
242
243        #[test]
244        fn test_pow() {
245            Python::with_gil(|py| {
246                let l = PyComplex::from_doubles_bound(py, 3.0, 1.2);
247                let r = PyComplex::from_doubles_bound(py, 1.2, 2.6);
248                let val = l.pow(&r);
249                assert_approx_eq!(val.real(), -1.419_309_997_016_603_7);
250                assert_approx_eq!(val.imag(), -0.541_297_466_033_544_6);
251            });
252        }
253    }
254}
255
256/// Implementation of functionality for [`PyComplex`].
257///
258/// These methods are defined for the `Bound<'py, PyComplex>` smart pointer, so to use method call
259/// syntax these methods are separated into a trait, because stable Rust does not yet support
260/// `arbitrary_self_types`.
261#[doc(alias = "PyComplex")]
262pub trait PyComplexMethods<'py>: crate::sealed::Sealed {
263    /// Returns the real part of the complex number.
264    fn real(&self) -> c_double;
265    /// Returns the imaginary part of the complex number.
266    fn imag(&self) -> c_double;
267    /// Returns `|self|`.
268    #[cfg(not(any(Py_LIMITED_API, PyPy, GraalPy)))]
269    fn abs(&self) -> c_double;
270    /// Returns `self` raised to the power of `other`.
271    #[cfg(not(any(Py_LIMITED_API, PyPy, GraalPy)))]
272    fn pow(&self, other: &Bound<'py, PyComplex>) -> Bound<'py, PyComplex>;
273}
274
275impl<'py> PyComplexMethods<'py> for Bound<'py, PyComplex> {
276    fn real(&self) -> c_double {
277        unsafe { ffi::PyComplex_RealAsDouble(self.as_ptr()) }
278    }
279
280    fn imag(&self) -> c_double {
281        unsafe { ffi::PyComplex_ImagAsDouble(self.as_ptr()) }
282    }
283
284    #[cfg(not(any(Py_LIMITED_API, PyPy, GraalPy)))]
285    fn abs(&self) -> c_double {
286        PyAnyMethods::abs(self.as_any())
287            .downcast_into()
288            .expect("Complex method __abs__ failed.")
289            .extract()
290            .expect("Failed to extract to c double.")
291    }
292
293    #[cfg(not(any(Py_LIMITED_API, PyPy, GraalPy)))]
294    fn pow(&self, other: &Bound<'py, PyComplex>) -> Bound<'py, PyComplex> {
295        Python::with_gil(|py| {
296            PyAnyMethods::pow(self.as_any(), other, py.None())
297                .downcast_into()
298                .expect("Complex method __pow__ failed.")
299        })
300    }
301}
302
303#[cfg(test)]
304mod tests {
305    use super::PyComplex;
306    use crate::{types::complex::PyComplexMethods, Python};
307    use assert_approx_eq::assert_approx_eq;
308
309    #[test]
310    fn test_from_double() {
311        use assert_approx_eq::assert_approx_eq;
312
313        Python::with_gil(|py| {
314            let complex = PyComplex::from_doubles_bound(py, 3.0, 1.2);
315            assert_approx_eq!(complex.real(), 3.0);
316            assert_approx_eq!(complex.imag(), 1.2);
317        });
318    }
319}
pyo3/types/complex.rs

pyo3/types/
complex.rs
