Skip to content

Commit cbbacba

Browse files
youknowoneyouknowone
committed
Refactor PyComplex
1 parent 801d011 commit cbbacba

1 file changed

Lines changed: 25 additions & 41 deletions

File tree

‎vm/src/obj/objcomplex.rs‎

Lines changed: 25 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ use num_traits::Zero;
33
use std::num::Wrapping;
44

55
use super::objfloat::{self, IntoPyFloat};
6-
use super::objtype::{self, PyClassRef};
6+
use super::objtype::PyClassRef;
77
use crate::function::OptionalArg;
88
use crate::pyhash;
99
use crate::pyobject::{
@@ -43,18 +43,15 @@ pub fn init(context: &PyContext) {
4343
PyComplex::extend_class(context, &context.types.complex_type);
4444
}
4545

46-
pub fn get_value(obj: &PyObjectRef) -> Complex64 {
47-
obj.payload::<PyComplex>().unwrap().value
48-
}
49-
5046
fn try_complex(value: &PyObjectRef, vm: &VirtualMachine) -> PyResult<Option<Complex64>> {
51-
Ok(if objtype::isinstance(&value, &vm.ctx.complex_type()) {
52-
Some(get_value(&value))
47+
let r = if let Some(complex) = value.payload_if_subclass::<PyComplex>(vm) {
48+
Some(complex.value)
5349
} else if let Some(float) = objfloat::try_float(value, vm)? {
5450
Some(Complex64::new(float, 0.0))
5551
} else {
5652
None
57-
})
53+
};
54+
Ok(r)
5855
}
5956

6057
#[pyimpl]
@@ -75,33 +72,35 @@ impl PyComplex {
7572
re.hypot(im)
7673
}
7774

78-
#[pymethod(name = "__add__")]
79-
fn add(&self, other: PyObjectRef, vm: &VirtualMachine) -> PyResult {
75+
#[inline]
76+
fn op<F>(&self, other: PyObjectRef, op: F, vm: &VirtualMachine) -> PyResult
77+
where
78+
F: Fn(Complex64, Complex64) -> Complex64,
79+
{
8080
try_complex(&other, vm)?.map_or_else(
8181
|| Ok(vm.ctx.not_implemented()),
82-
|other| (self.value + other).into_pyobject(vm),
82+
|other| op(self.value, other).into_pyobject(vm),
8383
)
8484
}
8585

86+
#[pymethod(name = "__add__")]
87+
fn add(&self, other: PyObjectRef, vm: &VirtualMachine) -> PyResult {
88+
self.op(other, |a, b| a + b, vm)
89+
}
90+
8691
#[pymethod(name = "__radd__")]
8792
fn radd(&self, other: PyObjectRef, vm: &VirtualMachine) -> PyResult {
8893
self.add(other, vm)
8994
}
9095

9196
#[pymethod(name = "__sub__")]
9297
fn sub(&self, other: PyObjectRef, vm: &VirtualMachine) -> PyResult {
93-
try_complex(&other, vm)?.map_or_else(
94-
|| Ok(vm.ctx.not_implemented()),
95-
|other| (self.value - other).into_pyobject(vm),
96-
)
98+
self.op(other, |a, b| a - b, vm)
9799
}
98100

99101
#[pymethod(name = "__rsub__")]
100102
fn rsub(&self, other: PyObjectRef, vm: &VirtualMachine) -> PyResult {
101-
try_complex(&other, vm)?.map_or_else(
102-
|| Ok(vm.ctx.not_implemented()),
103-
|other| (other - self.value).into_pyobject(vm),
104-
)
103+
self.op(other, |a, b| b - a, vm)
105104
}
106105

107106
#[pymethod(name = "conjugate")]
@@ -111,8 +110,8 @@ impl PyComplex {
111110

112111
#[pymethod(name = "__eq__")]
113112
fn eq(&self, other: PyObjectRef, vm: &VirtualMachine) -> PyObjectRef {
114-
let result = if objtype::isinstance(&other, &vm.ctx.complex_type()) {
115-
self.value == get_value(&other)
113+
let result = if let Some(other) = other.payload_if_subclass::<PyComplex>(vm) {
114+
self.value == other.value
116115
} else {
117116
match objfloat::try_float(&other, vm) {
118117
Ok(Some(other)) => self.value.im == 0.0f64 && self.value.re == other,
@@ -136,10 +135,7 @@ impl PyComplex {
136135

137136
#[pymethod(name = "__mul__")]
138137
fn mul(&self, other: PyObjectRef, vm: &VirtualMachine) -> PyResult {
139-
try_complex(&other, vm)?.map_or_else(
140-
|| Ok(vm.ctx.not_implemented()),
141-
|other| (self.value * other).into_pyobject(vm),
142-
)
138+
self.op(other, |a, b| a * b, vm)
143139
}
144140

145141
#[pymethod(name = "__rmul__")]
@@ -149,18 +145,12 @@ impl PyComplex {
149145

150146
#[pymethod(name = "__truediv__")]
151147
fn truediv(&self, other: PyObjectRef, vm: &VirtualMachine) -> PyResult {
152-
try_complex(&other, vm)?.map_or_else(
153-
|| Ok(vm.ctx.not_implemented()),
154-
|other| (self.value / other).into_pyobject(vm),
155-
)
148+
self.op(other, |a, b| a / b, vm)
156149
}
157150

158151
#[pymethod(name = "__rtruediv__")]
159152
fn rtruediv(&self, other: PyObjectRef, vm: &VirtualMachine) -> PyResult {
160-
try_complex(&other, vm)?.map_or_else(
161-
|| Ok(vm.ctx.not_implemented()),
162-
|other| (other / self.value).into_pyobject(vm),
163-
)
153+
self.op(other, |a, b| b / a, vm)
164154
}
165155

166156
#[pymethod(name = "__mod__")]
@@ -210,18 +200,12 @@ impl PyComplex {
210200

211201
#[pymethod(name = "__pow__")]
212202
fn pow(&self, other: PyObjectRef, vm: &VirtualMachine) -> PyResult {
213-
try_complex(&other, vm)?.map_or_else(
214-
|| Ok(vm.ctx.not_implemented()),
215-
|other| (self.value.powc(other)).into_pyobject(vm),
216-
)
203+
self.op(other, |a, b| a.powc(b), vm)
217204
}
218205

219206
#[pymethod(name = "__rpow__")]
220207
fn rpow(&self, other: PyObjectRef, vm: &VirtualMachine) -> PyResult {
221-
try_complex(&other, vm)?.map_or_else(
222-
|| Ok(vm.ctx.not_implemented()),
223-
|other| (other.powc(self.value)).into_pyobject(vm),
224-
)
208+
self.op(other, |a, b| b.powc(a), vm)
225209
}
226210

227211
#[pymethod(name = "__bool__")]

0 commit comments

Comments
 (0)