Fix math.ceil() and math.floor() to fall back to __ceil__ and __floor__

methods (respectively).  With Keir Mierle.

2 files changed, 70 insertions, 6 deletions

diff --git a/Lib/test/test_math.py b/Lib/test/test_math.py
index c28902b..2d15f3a 100644
--- a/Lib/test/test_math.py
+++ b/Lib/test/test_math.py
@@ -58,6 +58,19 @@ class MathTests(unittest.TestCase):
         self.ftest('ceil(-1.0)', math.ceil(-1.0), -1)
         self.ftest('ceil(-1.5)', math.ceil(-1.5), -1)
 
+        class TestCeil:
+            def __ceil__(self):
+                return 42
+        class TestNoCeil:
+            pass
+        self.ftest('ceil(TestCeil())', math.ceil(TestCeil()), 42)
+        self.assertRaises(TypeError, math.ceil, TestNoCeil())
+
+        t = TestNoCeil()
+        t.__ceil__ = lambda *args: args
+        self.assertRaises(TypeError, math.ceil, t)
+        self.assertRaises(TypeError, math.ceil, t, 0)
+
     def testCos(self):
         self.assertRaises(TypeError, math.cos)
         self.ftest('cos(-pi/2)', math.cos(-math.pi/2), 0)
@@ -101,6 +114,19 @@ class MathTests(unittest.TestCase):
         self.ftest('floor(1.23e167)', math.floor(1.23e167), 1.23e167)
         self.ftest('floor(-1.23e167)', math.floor(-1.23e167), -1.23e167)
 
+        class TestFloor:
+            def __floor__(self):
+                return 42
+        class TestNoFloor:
+            pass
+        self.ftest('floor(TestFloor())', math.floor(TestFloor()), 42)
+        self.assertRaises(TypeError, math.floor, TestNoFloor())
+
+        t = TestNoFloor()
+        t.__floor__ = lambda *args: args
+        self.assertRaises(TypeError, math.floor, t)
+        self.assertRaises(TypeError, math.floor, t, 0)
+
     def testFmod(self):
         self.assertRaises(TypeError, math.fmod)
         self.ftest('fmod(10,1)', math.fmod(10,1), 0)
diff --git a/Modules/mathmodule.c b/Modules/mathmodule.c
index 3d5640c..864df23 100644
--- a/Modules/mathmodule.c
+++ b/Modules/mathmodule.c
@@ -107,9 +107,28 @@ FUNC1(atan, atan,
 FUNC2(atan2, atan2,
       "atan2(y, x)\n\nReturn the arc tangent (measured in radians) of y/x.\n"
       "Unlike atan(y/x), the signs of both x and y are considered.")
-FUNC1(ceil, ceil,
-      "ceil(x)\n\nReturn the ceiling of x as a float.\n"
-      "This is the smallest integral value >= x.")
+
+static PyObject * math_ceil(PyObject *self, PyObject *number) {
+	static PyObject *ceil_str = NULL;
+	PyObject *method;
+
+	if (ceil_str == NULL) {
+		ceil_str = PyUnicode_FromString("__ceil__");
+		if (ceil_str == NULL)
+			return NULL;
+	}
+
+	method = _PyType_Lookup(Py_Type(number), ceil_str);
+	if (method == NULL)
+		return math_1(number, ceil);
+	else
+		return PyObject_CallFunction(method, "O", number);
+}
+
+PyDoc_STRVAR(math_ceil_doc,
+	     "ceil(x)\n\nReturn the ceiling of x as a float.\n"
+	     "This is the smallest integral value >= x.");
+
 FUNC1(cos, cos,
       "cos(x)\n\nReturn the cosine of x (measured in radians).")
 FUNC1(cosh, cosh,
@@ -118,9 +137,28 @@ FUNC1(exp, exp,
       "exp(x)\n\nReturn e raised to the power of x.")
 FUNC1(fabs, fabs,
       "fabs(x)\n\nReturn the absolute value of the float x.")
-FUNC1(floor, floor,
-      "floor(x)\n\nReturn the floor of x as a float.\n"
-      "This is the largest integral value <= x.")
+
+static PyObject * math_floor(PyObject *self, PyObject *number) {
+	static PyObject *floor_str = NULL;
+	PyObject *method;
+
+	if (floor_str == NULL) {
+		floor_str = PyUnicode_FromString("__floor__");
+		if (floor_str == NULL)
+			return NULL;
+	}
+
+	method = _PyType_Lookup(Py_Type(number), floor_str);
+	if (method == NULL)
+		return math_1(number, floor);
+	else
+		return PyObject_CallFunction(method, "O", number);
+}
+
+PyDoc_STRVAR(math_floor_doc,
+	     "floor(x)\n\nReturn the floor of x as a float.\n"
+	     "This is the largest integral value <= x.");
+
 FUNC2(fmod, fmod,
       "fmod(x,y)\n\nReturn fmod(x, y), according to platform C."
       "  x % y may differ.")