ENH: Add 'subok' parameter to PyArray_NewLikeArray, np.empty_like, np.zeros_like, and np.ones_like

This way, the sub-type can be avoided if necessary. This helps mitigate,
but doesn't fix, ticket #1753, by allowing "b = np.empty_like(a, subok=False)".

8 files changed, 108 insertions, 51 deletions

diff --git a/doc/source/reference/c-api.array.rst b/doc/source/reference/c-api.array.rst
index 05f077936..8c2b3a34e 100644
--- a/doc/source/reference/c-api.array.rst
+++ b/doc/source/reference/c-api.array.rst
@@ -158,7 +158,7 @@ From scratch
     *dims* and *strides* are copied into newly allocated dimension and
     strides arrays for the new array object.
 
-.. cfunction:: PyObject* PyArray_NewLikeArray(PyArrayObject* prototype, NPY_ORDER order, PyArray_Descr* descr)
+.. cfunction:: PyObject* PyArray_NewLikeArray(PyArrayObject* prototype, NPY_ORDER order, PyArray_Descr* descr, int subok)
 
     .. versionadded:: 1.6
 
@@ -176,6 +176,10 @@ From scratch
 
    If *descr* is NULL, the data type of *prototype* is used.
 
+   If *subok* is 1, the newly created array will use the sub-type of
+   *prototype* to create the new array, otherwise it will create a
+   base-class array.
+
 .. cfunction:: PyObject* PyArray_New(PyTypeObject* subtype, int nd, npy_intp* dims, int type_num, npy_intp* strides, void* data, int itemsize, int flags, PyObject* obj)
 
     This is similar to :cfunc:`PyArray_DescrNew` (...) except you
diff --git a/numpy/add_newdocs.py b/numpy/add_newdocs.py
index 8499b3d9d..5baf7424b 100644
--- a/numpy/add_newdocs.py
+++ b/numpy/add_newdocs.py
@@ -454,7 +454,7 @@ add_newdoc('numpy.core.multiarray', 'empty',
 
 add_newdoc('numpy.core.multiarray', 'empty_like',
     """
-    empty_like(a, dtype=None, order='K')
+    empty_like(a, dtype=None, order='K', subok=True)
 
     Return a new array with the same shape and type as a given array.
 
@@ -470,6 +470,10 @@ add_newdoc('numpy.core.multiarray', 'empty_like',
         'F' means F-order, 'A' means 'F' if ``a`` is Fortran contiguous,
         'C' otherwise. 'K' means match the layout of ``a`` as closely
         as possible.
+    subok : bool, optional.
+        If True, then the newly created array will use the sub-class
+        type of 'a', otherwise it will be a base-class array. Defaults
+        to True.
 
     Returns
     -------
diff --git a/numpy/core/numeric.py b/numpy/core/numeric.py
index 187296efe..dce617846 100644
--- a/numpy/core/numeric.py
+++ b/numpy/core/numeric.py
@@ -62,7 +62,7 @@ ufunc = type(sin)
 
 
 # originally from Fernando Perez's IPython
-def zeros_like(a, dtype=None, order='K'):
+def zeros_like(a, dtype=None, order='K', subok=True):
     """
     Return an array of zeros with the same shape and type as a given array.
 
@@ -112,7 +112,7 @@ def zeros_like(a, dtype=None, order='K'):
     array([ 0.,  0.,  0.])
 
     """
-    res = empty_like(a, dtype=dtype, order=order)
+    res = empty_like(a, dtype=dtype, order=order, subok=subok)
     res.fill(0)
     return res
 
diff --git a/numpy/core/src/multiarray/convert.c b/numpy/core/src/multiarray/convert.c
index 134350222..5e25ec1a1 100644
--- a/numpy/core/src/multiarray/convert.c
+++ b/numpy/core/src/multiarray/convert.c
@@ -454,7 +454,8 @@ PyArray_FillWithZero(PyArrayObject *a)
 NPY_NO_EXPORT PyObject *
 PyArray_NewCopy(PyArrayObject *m1, NPY_ORDER order)
 {
-    PyArrayObject *ret = (PyArrayObject *)PyArray_NewLikeArray(m1, order, NULL);
+    PyArrayObject *ret = (PyArrayObject *)PyArray_NewLikeArray(
+                                                    m1, order, NULL, 1);
     if (ret == NULL) {
         return NULL;
     }
diff --git a/numpy/core/src/multiarray/ctors.c b/numpy/core/src/multiarray/ctors.c
index b137b34f4..8cfe5da51 100644
--- a/numpy/core/src/multiarray/ctors.c
+++ b/numpy/core/src/multiarray/ctors.c
@@ -492,7 +492,8 @@ PyArray_MoveInto(PyArrayObject *dst, PyArrayObject *src)
         /*
          * Allocate a temporary copy array.
          */
-        tmp = (PyArrayObject *)PyArray_NewLikeArray(dst, NPY_KEEPORDER, NULL);
+        tmp = (PyArrayObject *)PyArray_NewLikeArray(dst,
+                                        NPY_KEEPORDER, NULL, 0);
         if (tmp == NULL) {
             return -1;
         }
@@ -1135,12 +1136,14 @@ PyArray_NewFromDescr(PyTypeObject *subtype, PyArray_Descr *descr, int nd,
  *             NPY_ANYORDER - Fortran if prototype is Fortran, C otherwise.
  *             NPY_KEEPORDER - Keeps the axis ordering of prototype.
  * dtype     - If not NULL, overrides the data type of the result.
+ * subok     - If 1, use the prototype's array subtype, otherwise
+ *             always create a base-class array.
  *
  * NOTE: If dtype is not NULL, steals the dtype reference.
  */
 NPY_NO_EXPORT PyObject *
 PyArray_NewLikeArray(PyArrayObject *prototype, NPY_ORDER order,
-                     PyArray_Descr *dtype)
+                     PyArray_Descr *dtype, int subok)
 {
     PyObject *ret = NULL;
     int ndim = PyArray_NDIM(prototype);
@@ -1173,14 +1176,14 @@ PyArray_NewLikeArray(PyArrayObject *prototype, NPY_ORDER order,
 
     /* If it's not KEEPORDER, this is simple */
     if (order != NPY_KEEPORDER) {
-        ret = PyArray_NewFromDescr(Py_TYPE(prototype),
-                                            dtype,
-                                            ndim,
-                                            PyArray_DIMS(prototype),
-                                            NULL,
-                                            NULL,
-                                            order,
-                                            (PyObject *)prototype);
+        ret = PyArray_NewFromDescr(subok ? Py_TYPE(prototype) : &PyArray_Type,
+                                        dtype,
+                                        ndim,
+                                        PyArray_DIMS(prototype),
+                                        NULL,
+                                        NULL,
+                                        order,
+                                        subok ? (PyObject *)prototype : NULL);
     }
     /* KEEPORDER needs some analysis of the strides */
     else {
@@ -1200,14 +1203,14 @@ PyArray_NewLikeArray(PyArrayObject *prototype, NPY_ORDER order,
         }
 
         /* Finally, allocate the array */
-        ret = PyArray_NewFromDescr(Py_TYPE(prototype),
-                                            dtype,
-                                            ndim,
-                                            shape,
-                                            strides,
-                                            NULL,
-                                            0,
-                                            (PyObject *)prototype);
+        ret = PyArray_NewFromDescr( subok ? Py_TYPE(prototype) : &PyArray_Type,
+                                        dtype,
+                                        ndim,
+                                        shape,
+                                        strides,
+                                        NULL,
+                                        0,
+                                        subok ? (PyObject *)prototype : NULL);
     }
 
     return ret;
diff --git a/numpy/core/src/multiarray/multiarraymodule.c b/numpy/core/src/multiarray/multiarraymodule.c
index ba31b0be7..28c596d9c 100644
--- a/numpy/core/src/multiarray/multiarraymodule.c
+++ b/numpy/core/src/multiarray/multiarraymodule.c
@@ -1730,20 +1730,22 @@ static PyObject *
 array_empty_like(PyObject *NPY_UNUSED(ignored), PyObject *args, PyObject *kwds)
 {
 
-    static char *kwlist[] = {"prototype","dtype","order",NULL};
+    static char *kwlist[] = {"prototype","dtype","order","subok",NULL};
     PyArrayObject *prototype = NULL;
     PyArray_Descr *dtype = NULL;
     NPY_ORDER order = NPY_KEEPORDER;
     PyObject *ret = NULL;
+    int subok = 1;
 
-    if (!PyArg_ParseTupleAndKeywords(args, kwds, "O&|O&O&", kwlist,
+    if (!PyArg_ParseTupleAndKeywords(args, kwds, "O&|O&O&i", kwlist,
                 PyArray_Converter, &prototype,
                 PyArray_DescrConverter2, &dtype,
-                PyArray_OrderConverter, &order)) {
+                PyArray_OrderConverter, &order,
+                &subok)) {
         goto fail;
     }
     /* steals the reference to dtype if it's not NULL */
-    ret = PyArray_NewLikeArray(prototype, order, dtype);
+    ret = PyArray_NewLikeArray(prototype, order, dtype, subok);
     Py_DECREF(prototype);
     return ret;
 
diff --git a/numpy/core/src/umath/ufunc_object.c b/numpy/core/src/umath/ufunc_object.c
index 1d53dd308..59722cef4 100644
--- a/numpy/core/src/umath/ufunc_object.c
+++ b/numpy/core/src/umath/ufunc_object.c
@@ -253,6 +253,16 @@ _find_array_prepare(PyObject *args, PyObject *kwds,
     PyObject *with_prep[NPY_MAXARGS], *preps[NPY_MAXARGS];
     PyObject *obj, *prep = NULL;
 
+    /* If a 'subok' parameter is passed and isn't True, don't wrap */
+    if (kwds != NULL && (obj = PyDict_GetItemString(kwds, "subok")) != NULL) {
+        if (obj != Py_True) {
+            for (i = 0; i < nout; i++) {
+                output_prep[i] = NULL;
+            }
+            return;
+        }
+    }
+
     nargs = PyTuple_GET_SIZE(args);
     for (i = 0; i < nin; i++) {
         obj = PyTuple_GET_ITEM(args, i);
@@ -706,6 +716,7 @@ static int get_ufunc_arguments(PyUFuncObject *self,
                 NPY_CASTING *out_casting,
                 PyObject **out_extobj,
                 PyObject **out_typetup,
+                int *out_subok,
                 int *out_any_object)
 {
     npy_intp i, nargs, nin = self->nin;
@@ -885,6 +896,15 @@ static int get_ufunc_arguments(PyUFuncObject *self,
                         Py_INCREF(value);
                         bad_arg = 0;
                     }
+                    else if (strncmp(str,"subok",5) == 0) {
+                        if (!PyBool_Check(value)) {
+                            PyErr_SetString(PyExc_TypeError,
+                                        "'subok' must be a boolean");
+                            goto fail;
+                        }
+                        *out_subok = (value == Py_True);
+                        bad_arg = 0;
+                    }
                     break;
                 case 'd':
                     /* Another way to specify 'sig' */
@@ -2056,7 +2076,7 @@ PyUFunc_GeneralizedFunction(PyUFuncObject *self,
     int nin, nout;
     int i, idim, niter;
     char *ufunc_name;
-    int retval = -1, any_object = 0;
+    int retval = -1, any_object = 0, subok = 1;
     NPY_CASTING input_casting;
 
     PyArray_Descr *dtype[NPY_MAXARGS];
@@ -2131,7 +2151,7 @@ PyUFunc_GeneralizedFunction(PyUFuncObject *self,
 
     /* Get all the arguments */
     retval = get_ufunc_arguments(self, args, kwds,
-                op, &order, &casting, &extobj, &type_tup, &any_object);
+                op, &order, &casting, &extobj, &type_tup, &subok, &any_object);
     if (retval < 0) {
         goto fail;
     }
@@ -2268,17 +2288,19 @@ PyUFunc_GeneralizedFunction(PyUFuncObject *self,
     printf("\n");
 #endif
 
-    /*
-     * Get the appropriate __array_prepare__ function to call
-     * for each output
-     */
-    _find_array_prepare(args, kwds, arr_prep, nin, nout);
+    if (subok) {
+        /*
+         * Get the appropriate __array_prepare__ function to call
+         * for each output
+         */
+        _find_array_prepare(args, kwds, arr_prep, nin, nout);
 
-    /* Set up arr_prep_args if a prep function was needed */
-    for (i = 0; i < nout; ++i) {
-        if (arr_prep[i] != NULL && arr_prep[i] != Py_None) {
-            arr_prep_args = make_arr_prep_args(nin, args, kwds);
-            break;
+        /* Set up arr_prep_args if a prep function was needed */
+        for (i = 0; i < nout; ++i) {
+            if (arr_prep[i] != NULL && arr_prep[i] != Py_None) {
+                arr_prep_args = make_arr_prep_args(nin, args, kwds);
+                break;
+            }
         }
     }
 
@@ -2472,7 +2494,7 @@ PyUFunc_GenericFunction(PyUFuncObject *self,
     int nin, nout;
     int i, niter;
     char *ufunc_name;
-    int retval = -1, any_object = 0;
+    int retval = -1, any_object = 0, subok = 1;
     NPY_CASTING input_casting;
 
     PyArray_Descr *dtype[NPY_MAXARGS];
@@ -2536,7 +2558,7 @@ PyUFunc_GenericFunction(PyUFuncObject *self,
 
     /* Get all the arguments */
     retval = get_ufunc_arguments(self, args, kwds,
-                op, &order, &casting, &extobj, &type_tup, &any_object);
+                op, &order, &casting, &extobj, &type_tup, &subok, &any_object);
     if (retval < 0) {
         goto fail;
     }
@@ -2614,17 +2636,19 @@ PyUFunc_GenericFunction(PyUFuncObject *self,
     printf("\n");
 #endif
 
-    /*
-     * Get the appropriate __array_prepare__ function to call
-     * for each output
-     */
-    _find_array_prepare(args, kwds, arr_prep, nin, nout);
+    if (subok) {
+        /*
+         * Get the appropriate __array_prepare__ function to call
+         * for each output
+         */
+        _find_array_prepare(args, kwds, arr_prep, nin, nout);
 
-    /* Set up arr_prep_args if a prep function was needed */
-    for (i = 0; i < nout; ++i) {
-        if (arr_prep[i] != NULL && arr_prep[i] != Py_None) {
-            arr_prep_args = make_arr_prep_args(nin, args, kwds);
-            break;
+        /* Set up arr_prep_args if a prep function was needed */
+        for (i = 0; i < nout; ++i) {
+            if (arr_prep[i] != NULL && arr_prep[i] != Py_None) {
+                arr_prep_args = make_arr_prep_args(nin, args, kwds);
+                break;
+            }
         }
     }
 
@@ -3967,6 +3991,16 @@ _find_array_wrap(PyObject *args, PyObject *kwds,
     PyObject *with_wrap[NPY_MAXARGS], *wraps[NPY_MAXARGS];
     PyObject *obj, *wrap = NULL;
 
+    /* If a 'subok' parameter is passed and isn't True, don't wrap */
+    if (kwds != NULL && (obj = PyDict_GetItemString(kwds, "subok")) != NULL) {
+        if (obj != Py_True) {
+            for (i = 0; i < nout; i++) {
+                output_wrap[i] = NULL;
+            }
+            return;
+        }
+    }
+
     nargs = PyTuple_GET_SIZE(args);
     for (i = 0; i < nin; i++) {
         obj = PyTuple_GET_ITEM(args, i);
diff --git a/numpy/core/tests/test_numeric.py b/numpy/core/tests/test_numeric.py
index 3752e0571..f7ecdfbcf 100644
--- a/numpy/core/tests/test_numeric.py
+++ b/numpy/core/tests/test_numeric.py
@@ -1195,6 +1195,15 @@ class TestLikeFuncs(TestCase):
             if not value is None:
                 assert_(all(dz == value))
 
+        # Test the 'subok' parameter'
+        a = np.matrix([[1,2],[3,4]])
+
+        b = like_function(a)
+        assert_(type(b) is np.matrix)
+
+        b = like_function(a, subok=False)
+        assert_(not (type(b) is np.matrix))
+
     def test_ones_like(self):
         self.check_like_function(np.ones_like, 1)