Merge pull request numpy#13371 from eric-wieser/__floor__-and-__ceil__

BUG/ENH: Make floor, ceil, and trunc call the matching special methods

doc/release/1.17.0-notes.rst

@@ -225,6 +225,12 @@ New keywords added to ``np.nan_to_num``
 user to define the value to replace the ``nan``, positive and negative ``np.inf`` values 
 respectively.
 
+`floor`, `ceil`, and `trunc` now respect builtin magic methods
+--------------------------------------------------------------
+These ufuncs now call the ``__floor__``, ``__ceil__``, and ``__trunc__``
+methods when called on object arrays, making them compatible with
+`decimal.Decimal` and `fractions.Fraction` objects.
+
 
 Changes
 =======

numpy/core/code_generators/generate_umath.py

@@ -765,14 +765,14 @@ def english_upper(s):
           docstrings.get('numpy.core.umath.ceil'),
           None,
           TD(flts, f='ceil', astype={'e':'f'}),
-          TD(P, f='ceil'),
+          TD(O, f='npy_ObjectCeil'),
           ),
 'trunc':
     Ufunc(1, 1, None,
           docstrings.get('numpy.core.umath.trunc'),
           None,
           TD(flts, f='trunc', astype={'e':'f'}),
-          TD(P, f='trunc'),
+          TD(O, f='npy_ObjectTrunc'),
           ),
 'fabs':
     Ufunc(1, 1, None,
@@ -786,7 +786,7 @@ def english_upper(s):
           docstrings.get('numpy.core.umath.floor'),
           None,
           TD(flts, f='floor', astype={'e':'f'}),
-          TD(P, f='floor'),
+          TD(O, f='npy_ObjectFloor'),
           ),
 'rint':
     Ufunc(1, 1, None,

numpy/core/src/umath/funcs.inc.src

@@ -159,6 +159,39 @@ npy_ObjectLogicalNot(PyObject *i1)
     }
 }
 
+static PyObject *
+npy_ObjectFloor(PyObject *obj) {
+    PyObject *math_floor_func = NULL;
+
+    npy_cache_import("math", "floor", &math_floor_func);
+    if (math_floor_func == NULL) {
+        return NULL;
+    }
+    return PyObject_CallFunction(math_floor_func, "O", obj);
+}
+
+static PyObject *
+npy_ObjectCeil(PyObject *obj) {
+    PyObject *math_ceil_func = NULL;
+
+    npy_cache_import("math", "ceil", &math_ceil_func);
+    if (math_ceil_func == NULL) {
+        return NULL;
+    }
+    return PyObject_CallFunction(math_ceil_func, "O", obj);
+}
+
+static PyObject *
+npy_ObjectTrunc(PyObject *obj) {
+    PyObject *math_trunc_func = NULL;
+
+    npy_cache_import("math", "trunc", &math_trunc_func);
+    if (math_trunc_func == NULL) {
+        return NULL;
+    }
+    return PyObject_CallFunction(math_trunc_func, "O", obj);
+}
+
 static PyObject *
 npy_ObjectGCD(PyObject *i1, PyObject *i2)
 {

numpy/core/tests/test_umath.py

@@ -5,6 +5,7 @@
 import fnmatch
 import itertools
 import pytest
+from fractions import Fraction
 
 import numpy.core.umath as ncu
 from numpy.core import _umath_tests as ncu_tests
@@ -2460,6 +2461,42 @@ def test_builtin_long(self):
         assert_equal(np.gcd(2**100, 3**100), 1)
 
 
+class TestRoundingFunctions(object):
+
+    def test_object_direct(self):
+        """ test direct implementation of these magic methods """
+        class C:
+            def __floor__(self):
+                return 1
+            def __ceil__(self):
+                return 2
+            def __trunc__(self):
+                return 3
+
+        arr = np.array([C(), C()])
+        assert_equal(np.floor(arr), [1, 1])
+        assert_equal(np.ceil(arr),  [2, 2])
+        assert_equal(np.trunc(arr), [3, 3])
+
+    def test_object_indirect(self):
+        """ test implementations via __float__ """
+        class C:
+            def __float__(self):
+                return -2.5
+
+        arr = np.array([C(), C()])
+        assert_equal(np.floor(arr), [-3, -3])
+        assert_equal(np.ceil(arr),  [-2, -2])
+        with pytest.raises(TypeError):
+            np.trunc(arr)  # consistent with math.trunc
+
+    def test_fraction(self):
+        f = Fraction(-4, 3)
+        assert_equal(np.floor(f), -2)
+        assert_equal(np.ceil(f), -1)
+        assert_equal(np.trunc(f), -1)
+
+
 class TestComplexFunctions(object):
     funcs = [np.arcsin,  np.arccos,  np.arctan, np.arcsinh, np.arccosh,
              np.arctanh, np.sin,     np.cos,    np.tan,     np.exp,