diff options
Diffstat (limited to 'numpy/core')
| -rw-r--r-- | numpy/core/src/multiarray/multiarraymodule.c | 24 | ||||
| -rw-r--r-- | numpy/core/tests/test_einsum.py | 11 |
2 files changed, 28 insertions, 7 deletions
diff --git a/numpy/core/src/multiarray/multiarraymodule.c b/numpy/core/src/multiarray/multiarraymodule.c index 6d323dbd8..896a3b07e 100644 --- a/numpy/core/src/multiarray/multiarraymodule.c +++ b/numpy/core/src/multiarray/multiarraymodule.c @@ -2652,21 +2652,31 @@ einsum_list_to_subscripts(PyObject *obj, char *subscripts, int subsize) /* Subscript */ else if (PyInt_Check(item) || PyLong_Check(item)) { long s = PyInt_AsLong(item); - if ( s < 0 || s > 2*26) { + npy_bool bad_input = 0; + + if (subindex + 1 >= subsize) { PyErr_SetString(PyExc_ValueError, - "subscript is not within the valid range [0, 52]"); + "subscripts list is too long"); Py_DECREF(obj); return -1; } - if (s < 26) { - subscripts[subindex++] = 'A' + s; + + if ( s < 0 ) { + bad_input = 1; + } + else if (s < 26) { + subscripts[subindex++] = 'A' + (char)s; + } + else if (s < 2*26) { + subscripts[subindex++] = 'a' + (char)s - 26; } else { - subscripts[subindex++] = 'a' + s; + bad_input = 1; } - if (subindex >= subsize) { + + if (bad_input) { PyErr_SetString(PyExc_ValueError, - "subscripts list is too long"); + "subscript is not within the valid range [0, 52)"); Py_DECREF(obj); return -1; } diff --git a/numpy/core/tests/test_einsum.py b/numpy/core/tests/test_einsum.py index 104dd1986..d92398456 100644 --- a/numpy/core/tests/test_einsum.py +++ b/numpy/core/tests/test_einsum.py @@ -596,6 +596,17 @@ class TestEinSum(object): [[[1, 3], [3, 9], [5, 15], [7, 21]], [[8, 16], [16, 32], [24, 48], [32, 64]]]) + def test_subscript_range(self): + # Issue #7741, make sure that all letters of Latin alphabet (both uppercase & lowercase) can be used + # when creating a subscript from arrays + a = np.ones((2, 3)) + b = np.ones((3, 4)) + np.einsum(a, [0, 20], b, [20, 2], [0, 2], optimize=False) + np.einsum(a, [0, 27], b, [27, 2], [0, 2], optimize=False) + np.einsum(a, [0, 51], b, [51, 2], [0, 2], optimize=False) + assert_raises(ValueError, lambda: np.einsum(a, [0, 52], b, [52, 2], [0, 2], optimize=False)) + assert_raises(ValueError, lambda: np.einsum(a, [-1, 5], b, [5, 2], [-1, 2], optimize=False)) + def test_einsum_broadcast(self): # Issue #2455 change in handling ellipsis # remove the 'middle broadcast' error |