summaryrefslogtreecommitdiff
path: root/numpy
diff options
context:
space:
mode:
Diffstat (limited to 'numpy')
-rw-r--r--numpy/core/src/multiarray/convert_datatype.c19
-rw-r--r--numpy/core/tests/test_multiarray.py4
-rw-r--r--numpy/lib/function_base.py30
-rw-r--r--numpy/lib/tests/test_function_base.py52
4 files changed, 98 insertions, 7 deletions
diff --git a/numpy/core/src/multiarray/convert_datatype.c b/numpy/core/src/multiarray/convert_datatype.c
index 3df764a48..33a706412 100644
--- a/numpy/core/src/multiarray/convert_datatype.c
+++ b/numpy/core/src/multiarray/convert_datatype.c
@@ -2028,7 +2028,6 @@ PyArray_ConvertToCommonType(PyObject *op, int *retn)
{
int i, n, allscalars = 0;
PyArrayObject **mps = NULL;
- PyObject *otmp;
PyArray_Descr *intype = NULL, *stype = NULL;
PyArray_Descr *newtype = NULL;
NPY_SCALARKIND scalarkind = NPY_NOSCALAR, intypekind = NPY_NOSCALAR;
@@ -2067,9 +2066,13 @@ PyArray_ConvertToCommonType(PyObject *op, int *retn)
}
for (i = 0; i < n; i++) {
- otmp = PySequence_GetItem(op, i);
+ PyObject *otmp = PySequence_GetItem(op, i);
+ if (otmp == NULL) {
+ goto fail;
+ }
if (!PyArray_CheckAnyScalar(otmp)) {
newtype = PyArray_DescrFromObject(otmp, intype);
+ Py_DECREF(otmp);
Py_XDECREF(intype);
if (newtype == NULL) {
goto fail;
@@ -2079,6 +2082,7 @@ PyArray_ConvertToCommonType(PyObject *op, int *retn)
}
else {
newtype = PyArray_DescrFromObject(otmp, stype);
+ Py_DECREF(otmp);
Py_XDECREF(stype);
if (newtype == NULL) {
goto fail;
@@ -2088,7 +2092,6 @@ PyArray_ConvertToCommonType(PyObject *op, int *retn)
mps[i] = (PyArrayObject *)Py_None;
Py_INCREF(Py_None);
}
- Py_XDECREF(otmp);
}
if (intype == NULL) {
/* all scalars */
@@ -2112,6 +2115,9 @@ PyArray_ConvertToCommonType(PyObject *op, int *retn)
newtype = PyArray_PromoteTypes(intype, stype);
Py_XDECREF(intype);
intype = newtype;
+ if (newtype == NULL) {
+ goto fail;
+ }
}
for (i = 0; i < n; i++) {
Py_XDECREF(mps[i]);
@@ -2123,8 +2129,9 @@ PyArray_ConvertToCommonType(PyObject *op, int *retn)
/* Make sure all arrays are actual array objects. */
for (i = 0; i < n; i++) {
int flags = NPY_ARRAY_CARRAY;
+ PyObject *otmp = PySequence_GetItem(op, i);
- if ((otmp = PySequence_GetItem(op, i)) == NULL) {
+ if (otmp == NULL) {
goto fail;
}
if (!allscalars && ((PyObject *)(mps[i]) == Py_None)) {
@@ -2133,8 +2140,8 @@ PyArray_ConvertToCommonType(PyObject *op, int *retn)
Py_DECREF(Py_None);
}
Py_INCREF(intype);
- mps[i] = (PyArrayObject*)
- PyArray_FromAny(otmp, intype, 0, 0, flags, NULL);
+ mps[i] = (PyArrayObject*)PyArray_FromAny(otmp, intype, 0, 0,
+ flags, NULL);
Py_DECREF(otmp);
if (mps[i] == NULL) {
goto fail;
diff --git a/numpy/core/tests/test_multiarray.py b/numpy/core/tests/test_multiarray.py
index 277f64a5d..6fbdb7c6c 100644
--- a/numpy/core/tests/test_multiarray.py
+++ b/numpy/core/tests/test_multiarray.py
@@ -1416,6 +1416,10 @@ class TestMethods(object):
A = ind.choose((x, y2))
assert_equal(A, [[2, 2, 3], [2, 2, 3]])
+ oned = np.ones(1)
+ # gh-12031, caused SEGFAULT
+ assert_raises(TypeError, oned.choose,np.void(0), [oned])
+
def test_prod(self):
ba = [1, 2, 10, 11, 6, 5, 4]
ba2 = [[1, 2, 3, 4], [5, 6, 7, 9], [10, 3, 4, 5]]
diff --git a/numpy/lib/function_base.py b/numpy/lib/function_base.py
index 2992e92bb..704e56514 100644
--- a/numpy/lib/function_base.py
+++ b/numpy/lib/function_base.py
@@ -1090,7 +1090,7 @@ def gradient(f, *varargs, **kwargs):
return outvals
-def diff(a, n=1, axis=-1):
+def diff(a, n=1, axis=-1, prepend=np._NoValue, append=np._NoValue):
"""
Calculate the n-th discrete difference along the given axis.
@@ -1108,6 +1108,12 @@ def diff(a, n=1, axis=-1):
axis : int, optional
The axis along which the difference is taken, default is the
last axis.
+ prepend, append : array_like, optional
+ Values to prepend or append to "a" along axis prior to
+ performing the difference. Scalar values are expanded to
+ arrays with length 1 in the direction of axis and the shape
+ of the input array in along all other axes. Otherwise the
+ dimension and shape must match "a" except along axis.
Returns
-------
@@ -1176,6 +1182,28 @@ def diff(a, n=1, axis=-1):
nd = a.ndim
axis = normalize_axis_index(axis, nd)
+ combined = []
+ if prepend is not np._NoValue:
+ prepend = np.asanyarray(prepend)
+ if prepend.ndim == 0:
+ shape = list(a.shape)
+ shape[axis] = 1
+ prepend = np.broadcast_to(prepend, tuple(shape))
+ combined.append(prepend)
+
+ combined.append(a)
+
+ if append is not np._NoValue:
+ append = np.asanyarray(append)
+ if append.ndim == 0:
+ shape = list(a.shape)
+ shape[axis] = 1
+ append = np.broadcast_to(append, tuple(shape))
+ combined.append(append)
+
+ if len(combined) > 1:
+ a = np.concatenate(combined, axis)
+
slice1 = [slice(None)] * nd
slice2 = [slice(None)] * nd
slice1[axis] = slice(1, None)
diff --git a/numpy/lib/tests/test_function_base.py b/numpy/lib/tests/test_function_base.py
index d5faed6ae..40cca1dbb 100644
--- a/numpy/lib/tests/test_function_base.py
+++ b/numpy/lib/tests/test_function_base.py
@@ -734,6 +734,58 @@ class TestDiff(object):
assert_array_equal(out3.mask, [[], [], [], [], []])
assert_(type(out3) is type(x))
+ def test_prepend(self):
+ x = np.arange(5) + 1
+ assert_array_equal(diff(x, prepend=0), np.ones(5))
+ assert_array_equal(diff(x, prepend=[0]), np.ones(5))
+ assert_array_equal(np.cumsum(np.diff(x, prepend=0)), x)
+ assert_array_equal(diff(x, prepend=[-1, 0]), np.ones(6))
+
+ x = np.arange(4).reshape(2, 2)
+ result = np.diff(x, axis=1, prepend=0)
+ expected = [[0, 1], [2, 1]]
+ assert_array_equal(result, expected)
+ result = np.diff(x, axis=1, prepend=[[0], [0]])
+ assert_array_equal(result, expected)
+
+ result = np.diff(x, axis=0, prepend=0)
+ expected = [[0, 1], [2, 2]]
+ assert_array_equal(result, expected)
+ result = np.diff(x, axis=0, prepend=[[0, 0]])
+ assert_array_equal(result, expected)
+
+ assert_raises(ValueError, np.diff, x, prepend=np.zeros((3,3)))
+
+ assert_raises(np.AxisError, diff, x, prepend=0, axis=3)
+
+ def test_append(self):
+ x = np.arange(5)
+ result = diff(x, append=0)
+ expected = [1, 1, 1, 1, -4]
+ assert_array_equal(result, expected)
+ result = diff(x, append=[0])
+ assert_array_equal(result, expected)
+ result = diff(x, append=[0, 2])
+ expected = expected + [2]
+ assert_array_equal(result, expected)
+
+ x = np.arange(4).reshape(2, 2)
+ result = np.diff(x, axis=1, append=0)
+ expected = [[1, -1], [1, -3]]
+ assert_array_equal(result, expected)
+ result = np.diff(x, axis=1, append=[[0], [0]])
+ assert_array_equal(result, expected)
+
+ result = np.diff(x, axis=0, append=0)
+ expected = [[2, 2], [-2, -3]]
+ assert_array_equal(result, expected)
+ result = np.diff(x, axis=0, append=[[0, 0]])
+ assert_array_equal(result, expected)
+
+ assert_raises(ValueError, np.diff, x, append=np.zeros((3,3)))
+
+ assert_raises(np.AxisError, diff, x, append=0, axis=3)
+
class TestDelete(object):
View Lorry Controller Status