fix tests

AMYPAD · Nov 22, 2021 · 7643ef1 · 7643ef1
1 parent ccc3a45
commit 7643ef1
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 3 deletions.
diff --git a/cuvec/include/pycuvec.cuh b/cuvec/include/pycuvec.cuh
@@ -235,6 +235,22 @@ template <class T> PyCuVec<T> *PyCuVec_deepcopy(PyCuVec<T> *other) {
 template <class T> PyCuVec<T> *asPyCuVec(PyObject *o) {
   if (!o || Py_None == o) return NULL;
   if (PyObject_HasAttrString(o, "cuvec")) {
+    // NumPy type checking
+    PyObject *dtype = PyObject_GetAttrString(o, "dtype");
+    PyObject *c = PyObject_GetAttrString(dtype, "char");
+    Py_XDECREF(dtype);
+    if (PyUnicode_Check(c)) {
+      char *npchr = (char *)PyUnicode_1BYTE_DATA(c);
+      if (*npchr != *cuvec::PyType<T>::npchr()) {
+        PyErr_Format(PyExc_TypeError,
+                     "cannot convert underlying dtype('%s') to requested dtype('%s')", npchr,
+                     cuvec::PyType<T>::npchr());
+        Py_DECREF(c);
+        return NULL;
+      }
+    }
+    Py_XDECREF(c);
+    // return cuvec
     o = PyObject_GetAttrString(o, "cuvec");
     if (!o) return NULL;
     Py_DECREF(o);

diff --git a/tests/test_pycuvec.py b/tests/test_pycuvec.py
@@ -158,7 +158,8 @@ def test_np_types():
     d = cu.zeros((1337, 42), 'd')
     cu.asarray(increment2d_f(f))
     cu.asarray(increment2d_f(f, f))
-    with raises((TypeError, SystemError)):
-        cu.asarray(increment2d_f(f, d))
-    with raises((TypeError, SystemError)):
+    with raises(TypeError):
         cu.asarray(increment2d_f(d))
+    with raises(SystemError):
+        # the TypeError is suppressed since a new output is generated
+        cu.asarray(increment2d_f(f, d))