Fix tests for Python 3.10 (ARROW-14312)

apache · Oct 13, 2021 · de75bb5 · de75bb5
1 parent de8a6c0
commit de75bb5
Show file tree

Hide file tree

Showing 4 changed files with 38 additions and 10 deletions.
diff --git a/ci/scripts/python_wheel_unix_test.sh b/ci/scripts/python_wheel_unix_test.sh
@@ -51,7 +51,7 @@ export PARQUET_TEST_DATA=${source_dir}/submodules/parquet-testing/data
 
 if [ "${INSTALL_PYARROW}" == "ON" ]; then
   # Install the built wheels
-  pip install ${source_dir}/python/repaired_wheels/*.whl
+  pip install --force-reinstall ${source_dir}/python/repaired_wheels/*.whl
 fi
 
 if [ "${CHECK_IMPORTS}" == "ON" ]; then

diff --git a/cpp/src/arrow/python/helpers.cc b/cpp/src/arrow/python/helpers.cc
@@ -162,12 +162,41 @@ Status IntegerOverflowStatus(PyObject* obj, const std::string& overflow_message)
   }
 }
 
+Result<OwnedRef> PyObjectToPyInt(PyObject* obj) {
+  // Try to call __index__ or __int__ on `obj`
+  // (starting from Python 3.10, the latter isn't done anymore by PyLong_AsLong*).
+  OwnedRef ref(PyNumber_Index(obj));
+  if (ref) {
+    return ref;
+  }
+  PyErr_Clear();
+  const auto nb = Py_TYPE(obj)->tp_as_number;
+  if (nb && nb->nb_int) {
+    ref.reset(nb->nb_int(obj));
+    if (!ref) {
+      RETURN_IF_PYERROR();
+    }
+    DCHECK(ref);
+    return ref;
+  }
+  return Status::TypeError(
+      "object of type ",
+      PyObject_StdStringRepr(reinterpret_cast<PyObject*>(Py_TYPE(obj))),
+      " cannot be converted to int");
+}
+
 // Extract C signed int from Python object
 template <typename Int, enable_if_t<std::is_signed<Int>::value, Int> = 0>
 Status CIntFromPythonImpl(PyObject* obj, Int* out, const std::string& overflow_message) {
   static_assert(sizeof(Int) <= sizeof(long long),  // NOLINT
                 "integer type larger than long long");
 
+  OwnedRef ref;
+  if (!PyLong_Check(obj)) {
+    ARROW_ASSIGN_OR_RAISE(ref, PyObjectToPyInt(obj));
+    obj = ref.obj();
+  }
+
   if (sizeof(Int) > sizeof(long)) {  // NOLINT
     const auto value = PyLong_AsLongLong(obj);
     if (ARROW_PREDICT_FALSE(value == -1)) {
@@ -199,15 +228,11 @@ Status CIntFromPythonImpl(PyObject* obj, Int* out, const std::string& overflow_m
                 "integer type larger than unsigned long long");
 
   OwnedRef ref;
-  // PyLong_AsUnsignedLong() and PyLong_AsUnsignedLongLong() don't handle
-  // conversion from non-ints (e.g. np.uint64), so do it ourselves
   if (!PyLong_Check(obj)) {
-    ref.reset(PyNumber_Index(obj));
-    if (!ref) {
-      RETURN_IF_PYERROR();
-    }
+    ARROW_ASSIGN_OR_RAISE(ref, PyObjectToPyInt(obj));
     obj = ref.obj();
   }
+
   if (sizeof(Int) > sizeof(unsigned long)) {  // NOLINT
     const auto value = PyLong_AsUnsignedLongLong(obj);
     if (ARROW_PREDICT_FALSE(value == static_cast<decltype(value)>(-1))) {

diff --git a/cpp/src/arrow/python/python_to_arrow.cc b/cpp/src/arrow/python/python_to_arrow.cc
@@ -206,14 +206,16 @@ class PyValue {
   }
 
   template <typename T>
-  static enable_if_integer<T, Result<typename T::c_type>> Convert(const T*, const O&,
+  static enable_if_integer<T, Result<typename T::c_type>> Convert(const T* type, const O&,
                                                                   I obj) {
     typename T::c_type value;
     auto status = internal::CIntFromPython(obj, &value);
     if (ARROW_PREDICT_TRUE(status.ok())) {
       return value;
     } else if (!internal::PyIntScalar_Check(obj)) {
-      return internal::InvalidValue(obj, "tried to convert to int");
+      std::stringstream ss;
+      ss << "tried to convert to " << type->ToString();
+      return internal::InvalidValue(obj, ss.str());
     } else {
       return status;
     }

diff --git a/python/pyarrow/error.pxi b/python/pyarrow/error.pxi
@@ -175,7 +175,8 @@ def enable_signal_handlers(c_bool enable):
 
 # Whether we need a workaround for https://bugs.python.org/issue42248
 have_signal_refcycle = (sys.version_info < (3, 8, 10) or
-                        (3, 9) <= sys.version_info < (3, 9, 5))
+                        (3, 9) <= sys.version_info < (3, 9, 5) or
+                        sys.version_info[:2] == (3, 10))
 
 cdef class SignalStopHandler:
     cdef: