Move try_complex into PyObjectRef

This commit is contained in:
jfh
2021-10-07 12:23:12 +03:00
parent 7387324f59
commit acfdf9ceaa
2 changed files with 34 additions and 38 deletions

View File

@@ -39,6 +39,36 @@ impl From<Complex64> for PyComplex {
}
}
impl PyObjectRef {
/// Tries converting a python object into a complex, returns an option of whether the complex
/// and whether the object was a complex originally or coereced into one
pub fn try_complex(&self, vm: &VirtualMachine) -> PyResult<Option<(Complex64, bool)>> {
if let Some(complex) = self.payload_if_exact::<PyComplex>(vm) {
return Ok(Some((complex.value, true)));
}
if let Some(method) = vm.get_method(self.clone(), "__complex__") {
let result = vm.invoke(&method?, ())?;
// TODO: returning strict subclasses of complex in __complex__ is deprecated
return match result.payload::<PyComplex>() {
Some(complex_obj) => Ok(Some((complex_obj.value, true))),
None => Err(vm.new_type_error(format!(
"__complex__ returned non-complex (type '{}')",
result.class().name()
))),
};
}
// `complex` does not have a `__complex__` by default, so subclasses might not either,
// use the actual stored value in this case
if let Some(complex) = self.payload_if_subclass::<PyComplex>(vm) {
return Ok(Some((complex.value, true)));
}
if let Some(float) = self.try_to_f64(vm)? {
return Ok(Some((Complex64::new(float, 0.0), false)));
}
Ok(None)
}
}
pub fn init(context: &PyContext) {
PyComplex::extend_class(context, &context.types.complex_type);
}
@@ -99,7 +129,7 @@ impl SlotConstructor for PyComplex {
val
};
if let Some(c) = try_complex(&val, vm)? {
if let Some(c) = val.try_complex(vm)? {
c
} else if let Some(s) = val.payload_if_subclass::<PyStr>(vm) {
if args.imag.is_present() {
@@ -125,7 +155,7 @@ impl SlotConstructor for PyComplex {
// if an imaginary argument is not passed in
OptionalArg::Missing => (Complex64::new(real.im, 0.0), false),
OptionalArg::Present(obj) => {
if let Some(c) = try_complex(&obj, vm)? {
if let Some(c) = obj.try_complex(vm)? {
c
} else if obj.class().issubclass(&vm.ctx.types.str_type) {
return Err(
@@ -420,34 +450,3 @@ fn parse_str(s: &str) -> Option<Complex64> {
};
Some(value)
}
/// Tries converting a python object into a complex, returns an option of whether the complex
/// and whether the object was a complex originally or coereced into one
pub(crate) fn try_complex(
obj: &PyObjectRef,
vm: &VirtualMachine,
) -> PyResult<Option<(Complex64, bool)>> {
if let Some(complex) = obj.payload_if_exact::<PyComplex>(vm) {
return Ok(Some((complex.value, true)));
}
if let Some(method) = vm.get_method(obj.clone(), "__complex__") {
let result = vm.invoke(&method?, ())?;
// TODO: returning strict subclasses of complex in __complex__ is deprecated
return match result.payload::<PyComplex>() {
Some(complex_obj) => Ok(Some((complex_obj.value, true))),
None => Err(vm.new_type_error(format!(
"__complex__ returned non-complex (type '{}')",
result.class().name()
))),
};
}
// `complex` does not have a `__complex__` by default, so subclasses might not either,
// use the actual stored value in this case
if let Some(complex) = obj.payload_if_subclass::<PyComplex>(vm) {
return Ok(Some((complex.value, true)));
}
if let Some(float) = obj.try_to_f64(vm)? {
return Ok(Some((Complex64::new(float, 0.0), false)));
}
Ok(None)
}

View File

@@ -1,7 +1,4 @@
use crate::{
builtins::complex::try_complex, PyObjectRef, PyResult, TryFromObject, TypeProtocol,
VirtualMachine,
};
use crate::{PyObjectRef, PyResult, TryFromObject, TypeProtocol, VirtualMachine};
use num_complex::Complex64;
#[derive(Debug, Copy, Clone, PartialEq)]
@@ -19,7 +16,7 @@ impl ArgComplexLike {
impl TryFromObject for ArgComplexLike {
fn try_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<Self> {
// We do not care if it was already a complex.
let (value, _) = try_complex(&obj, vm)?.ok_or_else(|| {
let (value, _) = obj.try_complex(vm)?.ok_or_else(|| {
vm.new_type_error(format!("must be real number, not {}", obj.class().name()))
})?;
Ok(ArgComplexLike { value })