implement `PyComplexMethods`

This commit is contained in:
Icxolu 2024-01-27 02:02:58 +01:00
parent 0973da27e9
commit 37e2a4d9c9
4 changed files with 197 additions and 85 deletions

View File

@ -170,8 +170,8 @@ impl Number {
self.0 as f64
}
fn __complex__<'py>(&self, py: Python<'py>) -> &'py PyComplex {
PyComplex::from_doubles(py, self.0 as f64, 0.0)
fn __complex__<'py>(&self, py: Python<'py>) -> Bound<'py, PyComplex> {
PyComplex::from_doubles_bound(py, self.0 as f64, 0.0)
}
}
```
@ -321,8 +321,8 @@ impl Number {
self.0 as f64
}
fn __complex__<'py>(&self, py: Python<'py>) -> &'py PyComplex {
PyComplex::from_doubles(py, self.0 as f64, 0.0)
fn __complex__<'py>(&self, py: Python<'py>) -> Bound<'py, PyComplex> {
PyComplex::from_doubles_bound(py, self.0 as f64, 0.0)
}
}

View File

@ -59,10 +59,10 @@
//! #
//! # module.add_function(wrap_pyfunction!(get_eigenvalues, module)?)?;
//! #
//! # let m11 = PyComplex::from_doubles(py, 0_f64, -1_f64);
//! # let m12 = PyComplex::from_doubles(py, 1_f64, 0_f64);
//! # let m21 = PyComplex::from_doubles(py, 2_f64, -1_f64);
//! # let m22 = PyComplex::from_doubles(py, -1_f64, 0_f64);
//! # let m11 = PyComplex::from_doubles_bound(py, 0_f64, -1_f64);
//! # let m12 = PyComplex::from_doubles_bound(py, 1_f64, 0_f64);
//! # let m21 = PyComplex::from_doubles_bound(py, 2_f64, -1_f64);
//! # let m22 = PyComplex::from_doubles_bound(py, -1_f64, 0_f64);
//! #
//! # let result = module
//! # .getattr("get_eigenvalues")?

View File

@ -281,8 +281,11 @@ impl Dummy {
slf
}
fn __complex__<'py>(&self, py: crate::Python<'py>) -> &'py crate::types::PyComplex {
crate::types::PyComplex::from_doubles(py, 0.0, 0.0)
fn __complex__<'py>(
&self,
py: crate::Python<'py>,
) -> crate::Bound<'py, crate::types::PyComplex> {
crate::types::PyComplex::from_doubles_bound(py, 0.0, 0.0)
}
fn __int__(&self) -> u32 {
@ -673,8 +676,11 @@ impl Dummy {
slf
}
fn __complex__<'py>(&self, py: crate::Python<'py>) -> &'py crate::types::PyComplex {
crate::types::PyComplex::from_doubles(py, 0.0, 0.0)
fn __complex__<'py>(
&self,
py: crate::Python<'py>,
) -> crate::Bound<'py, crate::types::PyComplex> {
crate::types::PyComplex::from_doubles_bound(py, 0.0, 0.0)
}
fn __int__(&self) -> u32 {

View File

@ -1,4 +1,4 @@
use crate::{ffi, PyAny, Python};
use crate::{ffi, types::any::PyAnyMethods, Bound, PyAny, PyNativeType, Python};
use std::os::raw::c_double;
/// Represents a Python [`complex`](https://docs.python.org/3/library/functions.html#complex) object.
@ -19,118 +19,173 @@ pyobject_native_type!(
);
impl PyComplex {
/// Creates a new `PyComplex` from the given real and imaginary values.
/// Deprecated form of [`PyComplex::from_doubles_bound`]
#[cfg_attr(
not(feature = "gil-refs"),
deprecated(
since = "0.21.0",
note = "`PyComplex::from_doubles` will be replaced by `PyComplex::from_doubles_bound` in a future PyO3 version"
)
)]
pub fn from_doubles(py: Python<'_>, real: c_double, imag: c_double) -> &PyComplex {
Self::from_doubles_bound(py, real, imag).into_gil_ref()
}
/// Creates a new `PyComplex` from the given real and imaginary values.
pub fn from_doubles_bound(
py: Python<'_>,
real: c_double,
imag: c_double,
) -> Bound<'_, PyComplex> {
use crate::ffi_ptr_ext::FfiPtrExt;
unsafe {
let ptr = ffi::PyComplex_FromDoubles(real, imag);
py.from_owned_ptr(ptr)
ffi::PyComplex_FromDoubles(real, imag)
.assume_owned(py)
.downcast_into_unchecked()
}
}
/// Returns the real part of the complex number.
pub fn real(&self) -> c_double {
unsafe { ffi::PyComplex_RealAsDouble(self.as_ptr()) }
self.as_borrowed().real()
}
/// Returns the imaginary part of the complex number.
pub fn imag(&self) -> c_double {
unsafe { ffi::PyComplex_ImagAsDouble(self.as_ptr()) }
self.as_borrowed().imag()
}
}
#[cfg(not(any(Py_LIMITED_API, PyPy)))]
mod not_limited_impls {
use crate::ffi_ptr_ext::FfiPtrExt;
use crate::Borrowed;
use super::*;
use std::ops::{Add, Div, Mul, Neg, Sub};
impl PyComplex {
/// Returns `|self|`.
pub fn abs(&self) -> c_double {
unsafe {
let val = (*(self.as_ptr() as *mut ffi::PyComplexObject)).cval;
ffi::_Py_c_abs(val)
}
self.as_borrowed().abs()
}
/// Returns `self` raised to the power of `other`.
pub fn pow(&self, other: &PyComplex) -> &PyComplex {
unsafe {
self.py()
.from_owned_ptr(complex_operation(self, other, ffi::_Py_c_pow))
}
pub fn pow<'py>(&'py self, other: &'py PyComplex) -> &'py PyComplex {
self.as_borrowed().pow(&other.as_borrowed()).into_gil_ref()
}
}
#[inline(always)]
unsafe fn complex_operation(
l: &PyComplex,
r: &PyComplex,
pub(super) unsafe fn complex_operation<'py>(
l: Borrowed<'_, 'py, PyComplex>,
r: Borrowed<'_, 'py, PyComplex>,
operation: unsafe extern "C" fn(ffi::Py_complex, ffi::Py_complex) -> ffi::Py_complex,
) -> *mut ffi::PyObject {
let l_val = (*(l.as_ptr() as *mut ffi::PyComplexObject)).cval;
let r_val = (*(r.as_ptr() as *mut ffi::PyComplexObject)).cval;
let l_val = (*l.as_ptr().cast::<ffi::PyComplexObject>()).cval;
let r_val = (*r.as_ptr().cast::<ffi::PyComplexObject>()).cval;
ffi::PyComplex_FromCComplex(operation(l_val, r_val))
}
impl<'py> Add for &'py PyComplex {
type Output = &'py PyComplex;
fn add(self, other: &'py PyComplex) -> &'py PyComplex {
unsafe {
self.py()
.from_owned_ptr(complex_operation(self, other, ffi::_Py_c_sum))
macro_rules! bin_ops {
($trait:ident, $fn:ident, $op:tt, $ffi:path) => {
impl<'py> $trait for Borrowed<'_, 'py, PyComplex> {
type Output = Bound<'py, PyComplex>;
fn $fn(self, other: Self) -> Self::Output {
unsafe {
complex_operation(self, other, $ffi)
.assume_owned(self.py())
.downcast_into_unchecked()
}
}
}
}
impl<'py> $trait for &'py PyComplex {
type Output = &'py PyComplex;
fn $fn(self, other: &'py PyComplex) -> &'py PyComplex {
(self.as_borrowed() $op other.as_borrowed()).into_gil_ref()
}
}
impl<'py> $trait for &Bound<'py, PyComplex> {
type Output = Bound<'py, PyComplex>;
fn $fn(self, other: &Bound<'py, PyComplex>) -> Bound<'py, PyComplex> {
self.as_borrowed() $op other.as_borrowed()
}
}
impl<'py> $trait<Bound<'py, PyComplex>> for &Bound<'py, PyComplex> {
type Output = Bound<'py, PyComplex>;
fn $fn(self, other: Bound<'py, PyComplex>) -> Bound<'py, PyComplex> {
self.as_borrowed() $op other.as_borrowed()
}
}
impl<'py> $trait for Bound<'py, PyComplex> {
type Output = Bound<'py, PyComplex>;
fn $fn(self, other: Bound<'py, PyComplex>) -> Bound<'py, PyComplex> {
self.as_borrowed() $op other.as_borrowed()
}
}
impl<'py> $trait<&Self> for Bound<'py, PyComplex> {
type Output = Bound<'py, PyComplex>;
fn $fn(self, other: &Bound<'py, PyComplex>) -> Bound<'py, PyComplex> {
self.as_borrowed() $op other.as_borrowed()
}
}
};
}
impl<'py> Sub for &'py PyComplex {
type Output = &'py PyComplex;
fn sub(self, other: &'py PyComplex) -> &'py PyComplex {
unsafe {
self.py()
.from_owned_ptr(complex_operation(self, other, ffi::_Py_c_diff))
}
}
}
impl<'py> Mul for &'py PyComplex {
type Output = &'py PyComplex;
fn mul(self, other: &'py PyComplex) -> &'py PyComplex {
unsafe {
self.py()
.from_owned_ptr(complex_operation(self, other, ffi::_Py_c_prod))
}
}
}
impl<'py> Div for &'py PyComplex {
type Output = &'py PyComplex;
fn div(self, other: &'py PyComplex) -> &'py PyComplex {
unsafe {
self.py()
.from_owned_ptr(complex_operation(self, other, ffi::_Py_c_quot))
}
}
}
bin_ops!(Add, add, +, ffi::_Py_c_sum);
bin_ops!(Sub, sub, -, ffi::_Py_c_diff);
bin_ops!(Mul, mul, *, ffi::_Py_c_prod);
bin_ops!(Div, div, /, ffi::_Py_c_quot);
impl<'py> Neg for &'py PyComplex {
type Output = &'py PyComplex;
fn neg(self) -> &'py PyComplex {
unsafe {
let val = (*(self.as_ptr() as *mut ffi::PyComplexObject)).cval;
let val = (*self.as_ptr().cast::<ffi::PyComplexObject>()).cval;
self.py()
.from_owned_ptr(ffi::PyComplex_FromCComplex(ffi::_Py_c_neg(val)))
}
}
}
impl<'py> Neg for Borrowed<'_, 'py, PyComplex> {
type Output = Bound<'py, PyComplex>;
fn neg(self) -> Self::Output {
unsafe {
let val = (*self.as_ptr().cast::<ffi::PyComplexObject>()).cval;
ffi::PyComplex_FromCComplex(ffi::_Py_c_neg(val))
.assume_owned(self.py())
.downcast_into_unchecked()
}
}
}
impl<'py> Neg for &Bound<'py, PyComplex> {
type Output = Bound<'py, PyComplex>;
fn neg(self) -> Bound<'py, PyComplex> {
-self.as_borrowed()
}
}
impl<'py> Neg for Bound<'py, PyComplex> {
type Output = Bound<'py, PyComplex>;
fn neg(self) -> Bound<'py, PyComplex> {
-self.as_borrowed()
}
}
#[cfg(test)]
mod tests {
use super::PyComplex;
use crate::Python;
use crate::{types::complex::PyComplexMethods, Python};
use assert_approx_eq::assert_approx_eq;
#[test]
fn test_add() {
Python::with_gil(|py| {
let l = PyComplex::from_doubles(py, 3.0, 1.2);
let r = PyComplex::from_doubles(py, 1.0, 2.6);
let l = PyComplex::from_doubles_bound(py, 3.0, 1.2);
let r = PyComplex::from_doubles_bound(py, 1.0, 2.6);
let res = l + r;
assert_approx_eq!(res.real(), 4.0);
assert_approx_eq!(res.imag(), 3.8);
@ -140,8 +195,8 @@ mod not_limited_impls {
#[test]
fn test_sub() {
Python::with_gil(|py| {
let l = PyComplex::from_doubles(py, 3.0, 1.2);
let r = PyComplex::from_doubles(py, 1.0, 2.6);
let l = PyComplex::from_doubles_bound(py, 3.0, 1.2);
let r = PyComplex::from_doubles_bound(py, 1.0, 2.6);
let res = l - r;
assert_approx_eq!(res.real(), 2.0);
assert_approx_eq!(res.imag(), -1.4);
@ -151,8 +206,8 @@ mod not_limited_impls {
#[test]
fn test_mul() {
Python::with_gil(|py| {
let l = PyComplex::from_doubles(py, 3.0, 1.2);
let r = PyComplex::from_doubles(py, 1.0, 2.6);
let l = PyComplex::from_doubles_bound(py, 3.0, 1.2);
let r = PyComplex::from_doubles_bound(py, 1.0, 2.6);
let res = l * r;
assert_approx_eq!(res.real(), -0.12);
assert_approx_eq!(res.imag(), 9.0);
@ -162,8 +217,8 @@ mod not_limited_impls {
#[test]
fn test_div() {
Python::with_gil(|py| {
let l = PyComplex::from_doubles(py, 3.0, 1.2);
let r = PyComplex::from_doubles(py, 1.0, 2.6);
let l = PyComplex::from_doubles_bound(py, 3.0, 1.2);
let r = PyComplex::from_doubles_bound(py, 1.0, 2.6);
let res = l / r;
assert_approx_eq!(res.real(), 0.788_659_793_814_432_9);
assert_approx_eq!(res.imag(), -0.850_515_463_917_525_7);
@ -173,7 +228,7 @@ mod not_limited_impls {
#[test]
fn test_neg() {
Python::with_gil(|py| {
let val = PyComplex::from_doubles(py, 3.0, 1.2);
let val = PyComplex::from_doubles_bound(py, 3.0, 1.2);
let res = -val;
assert_approx_eq!(res.real(), -3.0);
assert_approx_eq!(res.imag(), -1.2);
@ -183,7 +238,7 @@ mod not_limited_impls {
#[test]
fn test_abs() {
Python::with_gil(|py| {
let val = PyComplex::from_doubles(py, 3.0, 1.2);
let val = PyComplex::from_doubles_bound(py, 3.0, 1.2);
assert_approx_eq!(val.abs(), 3.231_098_884_280_702_2);
});
}
@ -191,9 +246,9 @@ mod not_limited_impls {
#[test]
fn test_pow() {
Python::with_gil(|py| {
let l = PyComplex::from_doubles(py, 3.0, 1.2);
let r = PyComplex::from_doubles(py, 1.2, 2.6);
let val = l.pow(r);
let l = PyComplex::from_doubles_bound(py, 3.0, 1.2);
let r = PyComplex::from_doubles_bound(py, 1.2, 2.6);
let val = l.pow(&r);
assert_approx_eq!(val.real(), -1.419_309_997_016_603_7);
assert_approx_eq!(val.imag(), -0.541_297_466_033_544_6);
});
@ -201,10 +256,61 @@ mod not_limited_impls {
}
}
/// Implementation of functionality for [`PyComplex`].
///
/// These methods are defined for the `Bound<'py, PyComplex>` smart pointer, so to use method call
/// syntax these methods are separated into a trait, because stable Rust does not yet support
/// `arbitrary_self_types`.
#[doc(alias = "PyComplex")]
pub trait PyComplexMethods<'py> {
/// Returns the real part of the complex number.
fn real(&self) -> c_double;
/// Returns the imaginary part of the complex number.
fn imag(&self) -> c_double;
/// Returns `|self|`.
#[cfg(not(any(Py_LIMITED_API, PyPy)))]
fn abs(&self) -> c_double;
/// Returns `self` raised to the power of `other`.
#[cfg(not(any(Py_LIMITED_API, PyPy)))]
fn pow(&self, other: &Bound<'py, PyComplex>) -> Bound<'py, PyComplex>;
}
impl<'py> PyComplexMethods<'py> for Bound<'py, PyComplex> {
fn real(&self) -> c_double {
unsafe { ffi::PyComplex_RealAsDouble(self.as_ptr()) }
}
fn imag(&self) -> c_double {
unsafe { ffi::PyComplex_ImagAsDouble(self.as_ptr()) }
}
#[cfg(not(any(Py_LIMITED_API, PyPy)))]
fn abs(&self) -> c_double {
unsafe {
let val = (*self.as_ptr().cast::<ffi::PyComplexObject>()).cval;
ffi::_Py_c_abs(val)
}
}
#[cfg(not(any(Py_LIMITED_API, PyPy)))]
fn pow(&self, other: &Bound<'py, PyComplex>) -> Bound<'py, PyComplex> {
use crate::ffi_ptr_ext::FfiPtrExt;
unsafe {
not_limited_impls::complex_operation(
self.as_borrowed(),
other.as_borrowed(),
ffi::_Py_c_pow,
)
.assume_owned(self.py())
.downcast_into_unchecked()
}
}
}
#[cfg(test)]
mod tests {
use super::PyComplex;
use crate::Python;
use crate::{types::complex::PyComplexMethods, Python};
use assert_approx_eq::assert_approx_eq;
#[test]
@ -212,7 +318,7 @@ mod tests {
use assert_approx_eq::assert_approx_eq;
Python::with_gil(|py| {
let complex = PyComplex::from_doubles(py, 3.0, 1.2);
let complex = PyComplex::from_doubles_bound(py, 3.0, 1.2);
assert_approx_eq!(complex.real(), 3.0);
assert_approx_eq!(complex.imag(), 1.2);
});