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 ffiBound<'_, PyAny>::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}
278