summaryrefslogtreecommitdiff
path: root/numpy
diff options
context:
space:
mode:
Diffstat (limited to 'numpy')
-rw-r--r--numpy/core/ma.py14
-rw-r--r--numpy/core/src/arraymethods.c116
-rw-r--r--numpy/core/src/scalartypes.inc.src10
3 files changed, 134 insertions, 6 deletions
diff --git a/numpy/core/ma.py b/numpy/core/ma.py
index 18a4e9a8a..5bd7ed7d2 100644
--- a/numpy/core/ma.py
+++ b/numpy/core/ma.py
@@ -1360,16 +1360,22 @@ array(data = %(data)s,
return self._data.dtype
dtype = property(fget=_get_dtype, doc="type of the array elements.")
- def item(self):
- "Return Python scalar if possible."
+ def item(self, *args):
+ "Return Python scalar if possible"
if self._mask is not nomask:
- m = fromnumeric.ravel(self._mask)
+ m = self._mask.item(*args)
try:
if m[0]:
return masked
except IndexError:
return masked
- return self._data.item()
+ return self._data.item(*args)
+
+ def itemset(self, *args):
+ "Set Python scalar into array"
+ item = args[-1]
+ args = args[:-1]
+ self[args] = item
def tolist(self, fill_value=None):
"Convert to list"
diff --git a/numpy/core/src/arraymethods.c b/numpy/core/src/arraymethods.c
index d98f0b00d..25da318ea 100644
--- a/numpy/core/src/arraymethods.c
+++ b/numpy/core/src/arraymethods.c
@@ -534,6 +534,120 @@ array_toscalar(PyArrayObject *self, PyObject *args) {
}
}
+static PyObject *
+array_setscalar(PyArrayObject *self, PyObject *args) {
+ int n, nd;
+ int ret = -1;
+ PyObject *obj;
+ n = PyTuple_GET_SIZE(args)-1;
+
+ if (n < 0) {
+ PyErr_SetString(PyExc_ValueError,
+ "setitem must have at least one argument");
+ return NULL;
+ }
+ obj = PyTuple_GET_ITEM(args, n);
+ if (n==0) {
+ if (self->nd == 0 || PyArray_SIZE(self) == 1) {
+ ret = self->descr->f->setitem(obj, self->data, self);
+ }
+ else {
+ PyErr_SetString(PyExc_ValueError,
+ "can only place a scalar for an "
+ " array of size 1");
+ return NULL;
+ }
+ }
+ else if (n != self->nd && (n > 1 || self->nd==0)) {
+ PyErr_SetString(PyExc_ValueError,
+ "incorrect number of indices for " \
+ "array");
+ return NULL;
+ }
+ else if (n==1) { /* allows for flat setting as well as 1-d case */
+ intp value, loc, index, factor;
+ intp factors[MAX_DIMS];
+ PyObject *indobj;
+
+ indobj = PyTuple_GET_ITEM(args, 0);
+ if (PyTuple_Check(indobj)) {
+ PyObject *res;
+ PyObject *newargs;
+ PyObject *tmp;
+ int i, nn;
+ nn = PyTuple_GET_SIZE(indobj);
+ newargs = PyTuple_New(nn+1);
+ Py_INCREF(obj);
+ for (i=0; i<nn; i++) {
+ tmp = PyTuple_GET_ITEM(indobj, i);
+ Py_INCREF(tmp);
+ PyTuple_SET_ITEM(newargs, i, tmp);
+ }
+ PyTuple_SET_ITEM(newargs, nn, obj);
+ /* Call with a converted set of arguments */
+ res = array_setscalar(self, newargs);
+ Py_DECREF(newargs);
+ return res;
+ }
+ value = PyArray_PyIntAsIntp(indobj);
+ if (error_converting(value)) {
+ PyErr_SetString(PyExc_ValueError, "invalid integer");
+ return NULL;
+ }
+ if (value >= PyArray_SIZE(self)) {
+ PyErr_SetString(PyExc_ValueError,
+ "index out of bounds");
+ return NULL;
+ }
+ if (self->nd == 1) {
+ ret = self->descr->f->setitem(obj, self->data + value,
+ self);
+ goto finish;
+ }
+ nd = self->nd;
+ factor = 1;
+ while (nd--) {
+ factors[nd] = factor;
+ factor *= self->dimensions[nd];
+ }
+ loc = 0;
+ for (nd=0; nd < self->nd; nd++) {
+ index = value / factors[nd];
+ value = value % factors[nd];
+ loc += self->strides[nd]*index;
+ }
+
+ ret = self->descr->f->setitem(obj, self->data + loc, self);
+ }
+ else {
+ intp loc, index[MAX_DIMS];
+ PyObject *tupargs;
+ tupargs = PyTuple_GetSlice(args, 1, n+1);
+ nd = PyArray_IntpFromSequence(tupargs, index, MAX_DIMS);
+ Py_DECREF(tupargs);
+ if (nd < n) return NULL;
+ loc = 0;
+ while (nd--) {
+ if (index[nd] < 0)
+ index[nd] += self->dimensions[nd];
+ if (index[nd] < 0 ||
+ index[nd] >= self->dimensions[nd]) {
+ PyErr_SetString(PyExc_ValueError,
+ "index out of bounds");
+ return NULL;
+ }
+ loc += self->strides[nd]*index[nd];
+ }
+ ret = self->descr->f->setitem(obj, self->data + loc, self);
+ }
+
+ finish:
+ if (ret < 0) return NULL;
+ Py_INCREF(Py_None);
+ return Py_None;
+}
+
+
static PyObject *
array_cast(PyArrayObject *self, PyObject *args)
@@ -1683,6 +1797,8 @@ static PyMethodDef array_methods[] = {
METH_VARARGS | METH_KEYWORDS, NULL},
{"item", (PyCFunction)array_toscalar,
METH_VARARGS, NULL},
+ {"itemset", (PyCFunction) array_setscalar,
+ METH_VARARGS, NULL},
{"max", (PyCFunction)array_max,
METH_VARARGS | METH_KEYWORDS, NULL},
{"mean", (PyCFunction)array_mean,
diff --git a/numpy/core/src/scalartypes.inc.src b/numpy/core/src/scalartypes.inc.src
index 470de4d53..056a887d6 100644
--- a/numpy/core/src/scalartypes.inc.src
+++ b/numpy/core/src/scalartypes.inc.src
@@ -827,7 +827,6 @@ gentype_interface_get(PyObject *self)
-
static PyObject *
gentype_typedescr_get(PyObject *self)
{
@@ -1068,6 +1067,13 @@ gentype_@name@(PyObject *self, PyObject *args)
/**end repeat**/
static PyObject *
+gentype_itemset(PyObject *self, PyObject *args)
+{
+ PyErr_SetString(PyExc_ValueError, "array-scalars are immutable");
+ return NULL;
+}
+
+static PyObject *
gentype_squeeze(PyObject *self, PyObject *args)
{
if (!PyArg_ParseTuple(args, "")) return NULL;
@@ -1281,13 +1287,13 @@ gentype_setflags(PyObject *self, PyObject *args, PyObject *kwds)
return Py_None;
}
-
/* need to fill in doc-strings for these methods on import -- copy from
array docstrings
*/
static PyMethodDef gentype_methods[] = {
{"tolist", (PyCFunction)gentype_tolist, 1, NULL},
{"item", (PyCFunction)gentype_item, METH_VARARGS, NULL},
+ {"itemset", (PyCFunction)gentype_itemset, METH_VARARGS, NULL},
{"tofile", (PyCFunction)gentype_tofile,
METH_VARARGS|METH_KEYWORDS, NULL},
{"tostring", (PyCFunction)gentype_tostring, METH_VARARGS, NULL},