Add support for python 3

1 files changed, 17 insertions, 24 deletions

diff --git a/numpy/core/src/umath/struct_ufunc_test.c.src b/numpy/core/src/umath/struct_ufunc_test.c.src
index 3e9e017cb..5b861a047 100644
--- a/numpy/core/src/umath/struct_ufunc_test.c.src
+++ b/numpy/core/src/umath/struct_ufunc_test.c.src
@@ -48,7 +48,7 @@ static char types[3] = {NPY_UINT64, NPY_UINT64, NPY_UINT64};
 
 static void *data[1] = {NULL};
 
-#if PY_VERSION_HEX >= 0x03000000
+#if defined(NPY_PY3K)
 static struct PyModuleDef moduledef = {
     PyModuleDef_HEAD_INIT,
     "struct_ufunc_test",
@@ -60,40 +60,31 @@ static struct PyModuleDef moduledef = {
     NULL,
     NULL
 };
+#endif
 
-PyObject *PyInit_struct_ufunc_test(void)
-{
-    PyObject *m, *add_triplet, *d;
-    m = PyModule_Create(&moduledef);
-    if (!m) {
-        return NULL;
-    }
-
-    import_array();
-    import_umath();
-
-    add_triplet = PyUFunc_FromFuncAndData(funcs, data, types, 1, 2, 1,
-                                    PyUFunc_None, "add_triplet",
-                                    "add_triplet_docstring", 0);
-
-    d = PyModule_GetDict(m);
-
-    PyDict_SetItemString(d, "add_triplet", add_triplet);
-    Py_DECREF(add_triplet);
-
-    return m;
-}
+#if defined(NPY_PY3K)
+PyMODINIT_FUNC PyInit_struct_ufunc_test(void)
 #else
 PyMODINIT_FUNC initstruct_ufunc_test(void)
+#endif
 {
     PyObject *m, *add_triplet, *d;
     PyObject *dtype_dict;
     PyArray_Descr *dtype;
     PyArray_Descr *dtypes[3];
 
+#if defined(NPY_PY3K)
+    m = PyModule_Create(&moduledef);
+#else
     m = Py_InitModule("struct_ufunc_test", StructUfuncTestMethods);
+#endif
+
     if (m == NULL) {
+#if defined(NPY_PY3K)
+        return NULL;
+#else
         return;
+#endif
     }
 
     import_array();
@@ -122,5 +113,7 @@ PyMODINIT_FUNC initstruct_ufunc_test(void)
 
     PyDict_SetItemString(d, "add_triplet", add_triplet);
     Py_DECREF(add_triplet);
-}
+#if defined(NPY_PY3K)
+    return m;
 #endif
+}