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

pyo3/types/
complex.rs
