correctly lookup __trunc__ and __floor__

3 files changed, 17 insertions, 20 deletions

diff --git a/Lib/test/test_descr.py b/Lib/test/test_descr.py
index 10820ab..297cc35 100644
--- a/Lib/test/test_descr.py
+++ b/Lib/test/test_descr.py
@@ -1,6 +1,7 @@
 import builtins
 import sys
 import types
+import math
 import unittest
 
 from copy import deepcopy
@@ -1578,6 +1579,8 @@ order (MRO) for bases """
             ("__exit__", run_context, swallow, set(), {"__enter__" : iden}),
             ("__complex__", complex, complex_num, set(), {}),
             ("__format__", format, format_impl, set(), {}),
+            ("__floor__", math.floor, zero, set(), {}),
+            ("__trunc__", math.trunc, zero, set(), {}),
             ]
 
         class Checker(object):
diff --git a/Misc/NEWS b/Misc/NEWS
index 9cf23c3..2671d85 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -1374,6 +1374,9 @@ Library
 Extension Modules
 -----------------
 
+- In the math module, correctly lookup __trunc__ and __floor__ as special
+  methods.
+
 - Issue #9005: Prevent utctimetuple() from producing year 0 or year
   10,000.  Prior to this change, timezone adjustment in utctimetuple()
   could produce tm_year value of 0 or 10,000.  Now an OverflowError is
diff --git a/Modules/mathmodule.c b/Modules/mathmodule.c
index ef6980f..2f656bb 100644
--- a/Modules/mathmodule.c
+++ b/Modules/mathmodule.c
@@ -883,17 +883,13 @@ static PyObject * math_floor(PyObject *self, PyObject *number) {
     static PyObject *floor_str = NULL;
     PyObject *method;
 
-    if (floor_str == NULL) {
-        floor_str = PyUnicode_InternFromString("__floor__");
-        if (floor_str == NULL)
+    method = _PyObject_LookupSpecial(number, "__floor__", &floor_str);
+    if (method == NULL) {
+        if (PyErr_Occurred())
             return NULL;
-    }
-
-    method = _PyType_Lookup(Py_TYPE(number), floor_str);
-    if (method == NULL)
         return math_1_to_int(number, floor, 0);
-    else
-        return PyObject_CallFunction(method, "O", number);
+    }
+    return PyObject_CallFunctionObjArgs(method, NULL);
 }
 
 PyDoc_STRVAR(math_floor_doc,
@@ -1427,20 +1423,15 @@ math_trunc(PyObject *self, PyObject *number)
             return NULL;
     }
 
-    if (trunc_str == NULL) {
-        trunc_str = PyUnicode_InternFromString("__trunc__");
-        if (trunc_str == NULL)
-            return NULL;
-    }
-
-    trunc = _PyType_Lookup(Py_TYPE(number), trunc_str);
+    trunc = _PyObject_LookupSpecial(number, "__trunc__", &trunc_str);
     if (trunc == NULL) {
-        PyErr_Format(PyExc_TypeError,
-                     "type %.100s doesn't define __trunc__ method",
-                     Py_TYPE(number)->tp_name);
+        if (!PyErr_Occurred())
+            PyErr_Format(PyExc_TypeError,
+                         "type %.100s doesn't define __trunc__ method",
+                         Py_TYPE(number)->tp_name);
         return NULL;
     }
-    return PyObject_CallFunctionObjArgs(trunc, number, NULL);
+    return PyObject_CallFunctionObjArgs(trunc, NULL);
 }
 
 PyDoc_STRVAR(math_trunc_doc,