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

pyo3/types/
complex.rs
