3K: random: make mtrand to import

2 files changed, 27 insertions, 2 deletions

diff --git a/numpy/random/mtrand/mtrand.pyx b/numpy/random/mtrand/mtrand.pyx
index e98414ae6..0340bef62 100644
--- a/numpy/random/mtrand/mtrand.pyx
+++ b/numpy/random/mtrand/mtrand.pyx
@@ -31,6 +31,9 @@ cdef extern from "math.h":
     double sin(double x)
     double cos(double x)
 
+cdef extern from "mtrand_py_helper.h":
+    object empty_py_bytes(unsigned long length, void **bytes)
+
 cdef extern from "randomkit.h":
 
     ctypedef struct rk_state:
@@ -883,8 +886,7 @@ cdef class RandomState:
 
         """
         cdef void *bytes
-        bytestring = PyString_FromStringAndSize(NULL, length)
-        bytes = PyString_AS_STRING(bytestring)
+        bytestring = empty_py_bytes(length, &bytes)
         rk_fill(bytes, length, self.internal_state)
         return bytestring
 
diff --git a/numpy/random/mtrand/mtrand_py_helper.h b/numpy/random/mtrand/mtrand_py_helper.h
new file mode 100644
index 000000000..2e7e4b7dc
--- /dev/null
+++ b/numpy/random/mtrand/mtrand_py_helper.h
@@ -0,0 +1,23 @@
+#ifndef _MTRAND_PY_HELPER_H_
+#define _MTRAND_PY_HELPER_H_
+
+#include <Python.h>
+
+static PyObject *empty_py_bytes(unsigned long length, void **bytes)
+{
+    PyObject *b;
+#if PY_MAJOR_VERSION >= 3
+    b = PyBytes_FromStringAndSize(NULL, length);
+    if (b) {
+        *bytes = PyBytes_AS_STRING(b);
+    }
+#else
+    b = PyString_FromStringAndSize(NULL, length);
+    if (b) {
+        *bytes = PyString_AS_STRING(b);
+    }
+#endif
+    return b;
+}
+
+#endif /* _MTRAND_PY_HELPER_H_ */