1use crate::{ffi, PyAny, Python};
2use std::os::raw::c_double;
3
4/// Represents a Python [`complex`](https://docs.python.org/3/library/functions.html#complex) object.
5///
6/// Note that `PyComplex` supports only basic operations. For advanced operations
7/// consider using [num-complex](https://docs.rs/num-complex)'s [`Complex`] type instead.
8/// This optional dependency can be activated with the `num-complex` feature flag.
9///
10/// [`Complex`]: https://docs.rs/num-complex/latest/num_complex/struct.Complex.html
11#[repr(transparent)]
12pub struct PyComplex(PyAny);
13
14pyobject_native_type!(
15 PyComplex,
16 ffi::PyComplexObject,
17 pyobject_native_static_type_object!(ffi::PyComplex_Type),
18 #checkfunction=ffi::PyComplex_Check
19);
20
21impl PyComplex {
22 /// Creates a new `PyComplex` from the given real and imaginary values.
23 pub fn from_doubles(py: Python<'_>, real: c_double, imag: c_double) -> &PyComplex {
24 unsafe {
25 let ptr: *mut PyObject = ffi::PyComplex_FromDoubles(real, imag);
26 py.from_owned_ptr(ptr)
27 }
28 }
29 /// Returns the real part of the complex number.
30 pub fn real(&self) -> c_double {
31 unsafe { ffi::PyComplex_RealAsDouble(self.as_ptr()) }
32 }
33 /// Returns the imaginary part of the complex number.
34 pub fn imag(&self) -> c_double {
35 unsafe { ffi::PyComplex_ImagAsDouble(self.as_ptr()) }
36 }
37}
38
39#[cfg(not(any(Py_LIMITED_API, PyPy)))]
40mod not_limited_impls {
41 use super::*;
42 use std::ops::{Add, Div, Mul, Neg, Sub};
43
44 impl PyComplex {
45 /// Returns `|self|`.
46 pub fn abs(&self) -> c_double {
47 unsafe {
48 let val = (*(self.as_ptr() as *mut ffi::PyComplexObject)).cval;
49 ffi::_Py_c_abs(val)
50 }
51 }
52 /// Returns `self` raised to the power of `other`.
53 pub fn pow(&self, other: &PyComplex) -> &PyComplex {
54 unsafe {
55 self.py()
56 .from_owned_ptr(complex_operation(self, other, ffi::_Py_c_pow))
57 }
58 }
59 }
60
61 #[inline(always)]
62 unsafe fn complex_operation(
63 l: &PyComplex,
64 r: &PyComplex,
65 operation: unsafe extern "C" fn(ffi::Py_complex, ffi::Py_complex) -> ffi::Py_complex,
66 ) -> *mut ffi::PyObject {
67 let l_val = (*(l.as_ptr() as *mut ffi::PyComplexObject)).cval;
68 let r_val = (*(r.as_ptr() as *mut ffi::PyComplexObject)).cval;
69 ffi::PyComplex_FromCComplex(operation(l_val, r_val))
70 }
71
72 impl<'py> Add for &'py PyComplex {
73 type Output = &'py PyComplex;
74 fn add(self, other: &'py PyComplex) -> &'py PyComplex {
75 unsafe {
76 self.py()
77 .from_owned_ptr(complex_operation(self, other, ffi::_Py_c_sum))
78 }
79 }
80 }
81
82 impl<'py> Sub for &'py PyComplex {
83 type Output = &'py PyComplex;
84 fn sub(self, other: &'py PyComplex) -> &'py PyComplex {
85 unsafe {
86 self.py()
87 .from_owned_ptr(complex_operation(self, other, ffi::_Py_c_diff))
88 }
89 }
90 }
91
92 impl<'py> Mul for &'py PyComplex {
93 type Output = &'py PyComplex;
94 fn mul(self, other: &'py PyComplex) -> &'py PyComplex {
95 unsafe {
96 self.py()
97 .from_owned_ptr(complex_operation(self, other, ffi::_Py_c_prod))
98 }
99 }
100 }
101
102 impl<'py> Div for &'py PyComplex {
103 type Output = &'py PyComplex;
104 fn div(self, other: &'py PyComplex) -> &'py PyComplex {
105 unsafe {
106 self.py()
107 .from_owned_ptr(complex_operation(self, other, ffi::_Py_c_quot))
108 }
109 }
110 }
111
112 impl<'py> Neg for &'py PyComplex {
113 type Output = &'py PyComplex;
114 fn neg(self) -> &'py PyComplex {
115 unsafe {
116 let val = (*(self.as_ptr() as *mut ffi::PyComplexObject)).cval;
117 self.py()
118 .from_owned_ptr(ffi::PyComplex_FromCComplex(ffi::_Py_c_neg(val)))
119 }
120 }
121 }
122
123 #[cfg(test)]
124 mod tests {
125 use super::PyComplex;
126 use crate::Python;
127 use assert_approx_eq::assert_approx_eq;
128
129 #[test]
130 fn test_add() {
131 Python::with_gil(|py| {
132 let l = PyComplex::from_doubles(py, 3.0, 1.2);
133 let r = PyComplex::from_doubles(py, 1.0, 2.6);
134 let res = l + r;
135 assert_approx_eq!(res.real(), 4.0);
136 assert_approx_eq!(res.imag(), 3.8);
137 });
138 }
139
140 #[test]
141 fn test_sub() {
142 Python::with_gil(|py| {
143 let l = PyComplex::from_doubles(py, 3.0, 1.2);
144 let r = PyComplex::from_doubles(py, 1.0, 2.6);
145 let res = l - r;
146 assert_approx_eq!(res.real(), 2.0);
147 assert_approx_eq!(res.imag(), -1.4);
148 });
149 }
150
151 #[test]
152 fn test_mul() {
153 Python::with_gil(|py| {
154 let l = PyComplex::from_doubles(py, 3.0, 1.2);
155 let r = PyComplex::from_doubles(py, 1.0, 2.6);
156 let res = l * r;
157 assert_approx_eq!(res.real(), -0.12);
158 assert_approx_eq!(res.imag(), 9.0);
159 });
160 }
161
162 #[test]
163 fn test_div() {
164 Python::with_gil(|py| {
165 let l = PyComplex::from_doubles(py, 3.0, 1.2);
166 let r = PyComplex::from_doubles(py, 1.0, 2.6);
167 let res = l / r;
168 assert_approx_eq!(res.real(), 0.788_659_793_814_432_9);
169 assert_approx_eq!(res.imag(), -0.850_515_463_917_525_7);
170 });
171 }
172
173 #[test]
174 fn test_neg() {
175 Python::with_gil(|py| {
176 let val = PyComplex::from_doubles(py, 3.0, 1.2);
177 let res = -val;
178 assert_approx_eq!(res.real(), -3.0);
179 assert_approx_eq!(res.imag(), -1.2);
180 });
181 }
182
183 #[test]
184 fn test_abs() {
185 Python::with_gil(|py| {
186 let val = PyComplex::from_doubles(py, 3.0, 1.2);
187 assert_approx_eq!(val.abs(), 3.231_098_884_280_702_2);
188 });
189 }
190
191 #[test]
192 fn test_pow() {
193 Python::with_gil(|py| {
194 let l = PyComplex::from_doubles(py, 3.0, 1.2);
195 let r = PyComplex::from_doubles(py, 1.2, 2.6);
196 let val = l.pow(r);
197 assert_approx_eq!(val.real(), -1.419_309_997_016_603_7);
198 assert_approx_eq!(val.imag(), -0.541_297_466_033_544_6);
199 });
200 }
201 }
202}
203
204#[cfg(test)]
205mod tests {
206 use super::PyComplex;
207 use crate::Python;
208 use assert_approx_eq::assert_approx_eq;
209
210 #[test]
211 fn test_from_double() {
212 use assert_approx_eq::assert_approx_eq;
213
214 Python::with_gil(|py| {
215 let complex = PyComplex::from_doubles(py, 3.0, 1.2);
216 assert_approx_eq!(complex.real(), 3.0);
217 assert_approx_eq!(complex.imag(), 1.2);
218 });
219 }
220}
221