diff --git a/crates/vm/src/vm/vm_ops.rs b/crates/vm/src/vm/vm_ops.rs index 2e2c11f81..404228608 100644 --- a/crates/vm/src/vm/vm_ops.rs +++ b/crates/vm/src/vm/vm_ops.rs @@ -4,7 +4,7 @@ use crate::{ Py, PyRef, builtins::{PyInt, PyStr, PyStrInterned, PyStrRef, PyType, PyUtf8Str}, object::{AsObject, PyObject, PyObjectRef, PyResult}, - protocol::{PyNumberBinaryOp, PyNumberTernaryOp}, + protocol::{PyNumberBinaryOp, PyNumberSlots, PyNumberTernaryOp, PyNumberUnaryFunc}, types::PyComparisonOp, }; use num_traits::ToPrimitive; @@ -541,22 +541,30 @@ impl VirtualMachine { Err(self.new_unsupported_bin_op_error(a, b, "*=")) } + fn unary_op( + &self, + a: &PyObject, + slot: impl FnOnce(&PyNumberSlots) -> Option, + op: &str, + ) -> PyResult { + let f = slot(&a.class().slots.as_number) + .ok_or_else(|| self.new_unsupported_unary_error(a, op))?; + f(a.number(), self) + } + + // PyNumber_Absolute pub fn _abs(&self, a: &PyObject) -> PyResult { - self.get_special_method(a, identifier!(self, __abs__))? - .ok_or_else(|| self.new_unsupported_unary_error(a, "abs()"))? - .invoke((), self) + self.unary_op(a, |s| s.absolute.load(), "abs()") } + // PyNumber_Positive pub fn _pos(&self, a: &PyObject) -> PyResult { - self.get_special_method(a, identifier!(self, __pos__))? - .ok_or_else(|| self.new_unsupported_unary_error(a, "unary +"))? - .invoke((), self) + self.unary_op(a, |s| s.positive.load(), "unary +") } + // PyNumber_Negative pub fn _neg(&self, a: &PyObject) -> PyResult { - self.get_special_method(a, identifier!(self, __neg__))? - .ok_or_else(|| self.new_unsupported_unary_error(a, "unary -"))? - .invoke((), self) + self.unary_op(a, |s| s.negative.load(), "unary -") } pub fn _invert(&self, a: &PyObject) -> PyResult { @@ -571,9 +579,7 @@ impl VirtualMachine { self, )?; } - self.get_special_method(a, identifier!(self, __invert__))? - .ok_or_else(|| self.new_unsupported_unary_error(a, "unary ~"))? - .invoke((), self) + self.unary_op(a, |s| s.invert.load(), "unary ~") } // PyObject_Format diff --git a/extra_tests/snippets/operator_unary.py b/extra_tests/snippets/operator_unary.py new file mode 100644 index 000000000..9a61e3b9d --- /dev/null +++ b/extra_tests/snippets/operator_unary.py @@ -0,0 +1,75 @@ +from testutils import assert_raises + + +class A: + def __neg__(self): + return "neg" + + def __pos__(self): + return "pos" + + def __abs__(self): + return "abs" + + def __invert__(self): + return "invert" + + +a = A() +assert -a == "neg" +assert +a == "pos" +assert abs(a) == "abs" +assert ~a == "invert" + + +class I(int): + def __neg__(self): + return "I.neg" + + +assert -I(3) == "I.neg" +assert +I(3) == 3 +assert abs(I(-3)) == 3 +assert ~I(3) == -4 + + +# Slots follow methods assigned or deleted after class creation. +class Late: + pass + + +with assert_raises(TypeError): + -Late() +Late.__neg__ = lambda self: "late" +assert -Late() == "late" +del Late.__neg__ +with assert_raises(TypeError): + -Late() + +del I.__neg__ +assert -I(3) == -3 + +# Only the type is consulted, not the instance. +i = Late() +i.__neg__ = lambda: "instance" +with assert_raises(TypeError): + -i + + +class Meta(type): + def __neg__(cls): + return cls.__name__ + + +class WithMeta(metaclass=Meta): + pass + + +assert -WithMeta == "WithMeta" + +with assert_raises(TypeError) as cm: + -"a" +assert str(cm.exception) == "bad operand type for unary -: 'str'" +with assert_raises(TypeError) as cm: + abs([]) +assert str(cm.exception) == "bad operand type for abs(): 'list'"