diff --git a/src/wrapt-stubs/__init__.pyi b/src/wrapt-stubs/__init__.pyi index 0cca7a97..771d3004 100644 --- a/src/wrapt-stubs/__init__.pyi +++ b/src/wrapt-stubs/__init__.pyi @@ -149,6 +149,9 @@ if sys.version_info >= (3, 10): def __bytes__(self) -> bytes: ... def __index__(self) -> int: ... def __round__(self, ndigits: int | None = ..., /) -> Any: ... + def __trunc__(self) -> Any: ... + def __floor__(self) -> Any: ... + def __ceil__(self) -> Any: ... # Unary arithmetic. def __neg__(self) -> Any: ... diff --git a/src/wrapt/_wrappers.c b/src/wrapt/_wrappers.c index ec1fdd21..24516e7f 100644 --- a/src/wrapt/_wrappers.c +++ b/src/wrapt/_wrappers.c @@ -2537,6 +2537,64 @@ static PyObject *WraptObjectProxy_round(WraptObjectProxyObject *self, /* ------------------------------------------------------------------------- */ +static PyObject *WraptObjectProxy_call_math_unary(WraptObjectProxyObject *self, + const char *name) +{ + PyObject *module = NULL; + PyObject *function = NULL; + PyObject *result = NULL; + + if (!self->wrapped) + { + if (raise_uninitialized_wrapper_error(self) == -1) + return NULL; + } + + module = PyImport_ImportModule("math"); + + if (!module) + return NULL; + + function = PyObject_GetAttrString(module, name); + + if (!function) + { + Py_DECREF(module); + return NULL; + } + + Py_DECREF(module); + + result = PyObject_CallFunctionObjArgs(function, self->wrapped, NULL); + + Py_DECREF(function); + + return result; +} + +/* ------------------------------------------------------------------------- */ + +static PyObject *WraptObjectProxy_trunc(WraptObjectProxyObject *self) +{ + return WraptObjectProxy_call_math_unary(self, "trunc"); +} + +/* ------------------------------------------------------------------------- */ + +static PyObject *WraptObjectProxy_floor(WraptObjectProxyObject *self) +{ + return WraptObjectProxy_call_math_unary(self, "floor"); +} + +/* ------------------------------------------------------------------------- */ + +static PyObject *WraptObjectProxy_ceil(WraptObjectProxyObject *self) +{ + return WraptObjectProxy_call_math_unary(self, "ceil"); +} + +/* ------------------------------------------------------------------------- */ + static PyObject *WraptObjectProxy_complex(WraptObjectProxyObject *self, PyObject *args) { @@ -3148,6 +3206,9 @@ static PyMethodDef WraptObjectProxy_methods[] = { {"__reversed__", (PyCFunction)WraptObjectProxy_reversed, METH_NOARGS, 0}, {"__round__", (PyCFunction)WraptObjectProxy_round, METH_VARARGS | METH_KEYWORDS, 0}, + {"__trunc__", (PyCFunction)WraptObjectProxy_trunc, METH_NOARGS, 0}, + {"__floor__", (PyCFunction)WraptObjectProxy_floor, METH_NOARGS, 0}, + {"__ceil__", (PyCFunction)WraptObjectProxy_ceil, METH_NOARGS, 0}, {"__complex__", (PyCFunction)WraptObjectProxy_complex, METH_NOARGS, 0}, {"__mro_entries__", (PyCFunction)WraptObjectProxy_mro_entries, METH_VARARGS | METH_KEYWORDS, 0}, diff --git a/src/wrapt/wrappers.py b/src/wrapt/wrappers.py index eccc544a..35d73c6a 100644 --- a/src/wrapt/wrappers.py +++ b/src/wrapt/wrappers.py @@ -1,6 +1,7 @@ """Core object proxy and function wrapper implementations.""" import inspect +import math import operator import sys import types @@ -261,6 +262,15 @@ def __reversed__(self): def __round__(self, ndigits=None): return round(self.__wrapped__, ndigits) + def __trunc__(self): + return math.trunc(self.__wrapped__) + + def __floor__(self): + return math.floor(self.__wrapped__) + + def __ceil__(self): + return math.ceil(self.__wrapped__) + def __mro_entries__(self, bases): if not isinstance(self.__wrapped__, type) and hasattr( self.__wrapped__, "__mro_entries__" diff --git a/tests/core/test_object_proxy.py b/tests/core/test_object_proxy.py index 0996dfff..acc1d4c9 100644 --- a/tests/core/test_object_proxy.py +++ b/tests/core/test_object_proxy.py @@ -1973,6 +1973,54 @@ def test_fractions_round(self): self.assertEqual(round(instance, 3), round(proxy, 3)) self.assertEqual(round(instance, ndigits=3), round(proxy, ndigits=3)) + def test_math_integral_protocol_methods(self): + import math + + class Instance: + def __trunc__(self): + return "truncated" + + def __floor__(self): + return "floored" + + def __ceil__(self): + return "ceiled" + + instance = Instance() + proxy = wrapt.ObjectProxy(instance) + + self.assertEqual(math.trunc(proxy), "truncated") + self.assertEqual(math.floor(proxy), "floored") + self.assertEqual(math.ceil(proxy), "ceiled") + + def test_float_trunc(self): + import math + + instance = 1.7 + proxy = wrapt.ObjectProxy(instance) + + self.assertEqual(math.trunc(instance), math.trunc(proxy)) + self.assertEqual(math.trunc(instance), 1) + + def test_negative_float_trunc(self): + import math + + instance = -1.7 + proxy = wrapt.ObjectProxy(instance) + + self.assertEqual(math.trunc(instance), math.trunc(proxy)) + self.assertEqual(math.trunc(instance), -1) + + def test_fractions_trunc(self): + import fractions + import math + + instance = fractions.Fraction("7/3") + proxy = wrapt.ObjectProxy(instance) + + self.assertEqual(math.trunc(instance), math.trunc(proxy)) + self.assertEqual(math.trunc(instance), 2) + class TestArgumentUnpacking(unittest.TestCase):