From 508118e0002d7cd5b96865d6691c99bf6d91286a Mon Sep 17 00:00:00 2001 From: Padraic Fanning Date: Sat, 2 Jul 2022 22:25:34 -0400 Subject: [PATCH 1/3] Implement Number protocol for PyComplex --- vm/src/builtins/complex.rs | 99 ++++++++++++++++++++++++++++++++++++-- 1 file changed, 96 insertions(+), 3 deletions(-) diff --git a/vm/src/builtins/complex.rs b/vm/src/builtins/complex.rs index 9e4f5af904..4ed47b93b7 100644 --- a/vm/src/builtins/complex.rs +++ b/vm/src/builtins/complex.rs @@ -1,14 +1,15 @@ use super::{float, PyStr, PyType, PyTypeRef}; use crate::{ class::PyClassImpl, - convert::ToPyObject, + convert::{ToPyObject, ToPyResult}, function::{ OptionalArg, OptionalOption, PyArithmeticValue::{self, *}, PyComparisonValue, }, identifier, - types::{Comparable, Constructor, Hashable, PyComparisonOp}, + protocol::{PyNumber, PyNumberMethods}, + types::{AsNumber, Comparable, Constructor, Hashable, PyComparisonOp}, AsObject, Context, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine, }; use num_complex::Complex64; @@ -203,7 +204,7 @@ impl PyComplex { } } -#[pyimpl(flags(BASETYPE), with(Comparable, Hashable, Constructor))] +#[pyimpl(flags(BASETYPE), with(Comparable, Hashable, Constructor, AsNumber))] impl PyComplex { #[pymethod(magic)] fn complex(zelf: PyRef, vm: &VirtualMachine) -> PyRef { @@ -419,6 +420,98 @@ impl Hashable for PyComplex { } } +impl AsNumber for PyComplex { + const AS_NUMBER: PyNumberMethods = PyNumberMethods { + add: Some(|number, other, vm| Self::number_complex_op(number, other, |a, b| a + b, vm)), + subtract: Some(|number, other, vm| { + Self::number_complex_op(number, other, |a, b| a - b, vm) + }), + multiply: Some(|number, other, vm| { + Self::number_complex_op(number, other, |a, b| a * b, vm) + }), + remainder: None, + divmod: None, + power: Some(|number, other, vm| Self::number_general_op(number, other, inner_pow, vm)), + negative: Some(|number, vm| { + let value = Self::number_downcast(number).value; + (-value).to_pyresult(vm) + }), + positive: Some(|number, vm| Self::number_complex(number, vm).to_pyresult(vm)), + absolute: Some(|number, vm| { + let value = Self::number_downcast(number).value; + value.norm().to_pyresult(vm) + }), + boolean: Some(|number, _vm| Ok(Self::number_downcast(number).value.is_zero())), + invert: None, + lshift: None, + rshift: None, + and: None, + xor: None, + or: None, + int: None, + float: None, + inplace_add: None, + inplace_subtract: None, + inplace_multiply: None, + inplace_remainder: None, + inplace_divmod: None, + inplace_power: None, + inplace_lshift: None, + inplace_rshift: None, + inplace_and: None, + inplace_xor: None, + inplace_or: None, + floor_divide: None, + true_divide: Some(|number, other, vm| { + Self::number_general_op(number, other, inner_div, vm) + }), + inplace_floor_divide: None, + inplace_true_divide: None, + index: None, + matrix_multiply: None, + inplace_matrix_multiply: None, + }; +} + +impl PyComplex { + fn number_general_op( + number: &PyNumber, + other: &PyObject, + op: F, + vm: &VirtualMachine, + ) -> PyResult + where + F: FnOnce(Complex64, Complex64, &VirtualMachine) -> R, + R: ToPyResult, + { + if let (Some(a), Some(b)) = (number.obj.payload::(), other.payload::()) { + op(a.value, b.value, vm).to_pyresult(vm) + } else { + Ok(vm.ctx.not_implemented()) + } + } + + fn number_complex_op( + number: &PyNumber, + other: &PyObject, + op: F, + vm: &VirtualMachine, + ) -> PyResult + where + F: FnOnce(Complex64, Complex64) -> Complex64, + { + Self::number_general_op(number, other, |a, b, _vm| op(a, b), vm) + } + + fn number_complex(number: &PyNumber, vm: &VirtualMachine) -> PyRef { + if let Some(zelf) = number.obj.downcast_ref_if_exact::(vm) { + zelf.to_owned() + } else { + vm.ctx.new_complex(Self::number_downcast(number).value) + } + } +} + #[derive(FromArgs)] pub struct ComplexArgs { #[pyarg(any, optional)] From 9de82747802f4476594b15221411887f582698fe Mon Sep 17 00:00:00 2001 From: Padraic Fanning Date: Sun, 3 Jul 2022 14:39:36 -0400 Subject: [PATCH 2/3] Unmark passing test --- Lib/test/test_itertools.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/Lib/test/test_itertools.py b/Lib/test/test_itertools.py index cd5b49d16d..fea9a96a00 100644 --- a/Lib/test/test_itertools.py +++ b/Lib/test/test_itertools.py @@ -533,8 +533,6 @@ def test_compress(self): next(testIntermediate) self.assertEqual(list(op(testIntermediate)), list(result2)) - # TODO: RUSTPYTHON - @unittest.expectedFailure def test_count(self): self.assertEqual(lzip('abc',count()), [('a', 0), ('b', 1), ('c', 2)]) self.assertEqual(lzip('abc',count(3)), [('a', 3), ('b', 4), ('c', 5)]) From de9a4d5b3fefc38b3c897adbb806513b3dd14d84 Mon Sep 17 00:00:00 2001 From: Padraic Fanning Date: Fri, 8 Jul 2022 20:11:21 -0400 Subject: [PATCH 3/3] Replace `None` instances with update syntax --- vm/src/builtins/complex.rs | 28 +--------------------------- 1 file changed, 1 insertion(+), 27 deletions(-) diff --git a/vm/src/builtins/complex.rs b/vm/src/builtins/complex.rs index 4ed47b93b7..8805c62471 100644 --- a/vm/src/builtins/complex.rs +++ b/vm/src/builtins/complex.rs @@ -429,8 +429,6 @@ impl AsNumber for PyComplex { multiply: Some(|number, other, vm| { Self::number_complex_op(number, other, |a, b| a * b, vm) }), - remainder: None, - divmod: None, power: Some(|number, other, vm| Self::number_general_op(number, other, inner_pow, vm)), negative: Some(|number, vm| { let value = Self::number_downcast(number).value; @@ -442,34 +440,10 @@ impl AsNumber for PyComplex { value.norm().to_pyresult(vm) }), boolean: Some(|number, _vm| Ok(Self::number_downcast(number).value.is_zero())), - invert: None, - lshift: None, - rshift: None, - and: None, - xor: None, - or: None, - int: None, - float: None, - inplace_add: None, - inplace_subtract: None, - inplace_multiply: None, - inplace_remainder: None, - inplace_divmod: None, - inplace_power: None, - inplace_lshift: None, - inplace_rshift: None, - inplace_and: None, - inplace_xor: None, - inplace_or: None, - floor_divide: None, true_divide: Some(|number, other, vm| { Self::number_general_op(number, other, inner_div, vm) }), - inplace_floor_divide: None, - inplace_true_divide: None, - index: None, - matrix_multiply: None, - inplace_matrix_multiply: None, + ..PyNumberMethods::NOT_IMPLEMENTED }; }