Fixed Numpy return value for Python 3

jfsantos · Vijay Vasudevan · commit 32ed124dc90b · 2016-01-14T11:15:50.000-08:00
Using PY_MAJOR_VERSION instead of NUMPY_IMPORT_ARRAY_RETVAL

Fixed SWIG interface for Python 3 support

Make SWIG forward preprocessor directives to generated code
diff --git a/tensorflow/python/client/tf_session.i b/tensorflow/python/client/tf_session.i
@@ -204,7 +204,11 @@ tensorflow::ImportNumpy();
 // The wrapped function TF_GetOpList returns a TF_Buffer pointer. This typemap
 // creates a Python string from the TF_Buffer and returns it.
 %typemap(out) TF_Buffer TF_GetOpList {
+%#if PY_MAJOR_VERSION < 3
   $result = PyString_FromStringAndSize(
+%#else
+  $result = PyUnicode_FromStringAndSize(
+%#endif
       reinterpret_cast<const char*>($1.data), $1.length);
 }
 
diff --git a/tensorflow/python/lib/core/py_func.cc b/tensorflow/python/lib/core/py_func.cc
@@ -24,6 +24,14 @@ limitations under the License.
 #include "tensorflow/core/platform/mutex.h"
 #include "tensorflow/core/platform/port.h"
 
+// Return type of import_array() changed between Python 2 and 3
+// NUMPY_IMPORT_ARRAY_RETVAL is NULL for Python 3
+#if PY_MAJOR_VERSION >= 3
+#define NUMPY_IMPORT_ARRAY_RETURN_TYPE int
+#else
+#define NUMPY_IMPORT_ARRAY_RETURN_TYPE void
+#endif
+
 namespace tensorflow {
 namespace {
 
@@ -39,7 +47,7 @@ PyObject* GetPyTrampoline() {
 }
 
 // Module initialization (mainly import numpy) if needed.
-void InitIfNeeded() {
+NUMPY_IMPORT_ARRAY_RETURN_TYPE InitIfNeeded() {
   mutex_lock l(mu);
   if (!initialized) {
     PyGILState_STATE py_threadstate;
