implemented newaxis support for zero rank arrays

1 files changed, 63 insertions, 2 deletions

diff --git a/numpy/core/src/arrayobject.c b/numpy/core/src/arrayobject.c
index 57926a4a2..3537cbe21 100644
--- a/numpy/core/src/arrayobject.c
+++ b/numpy/core/src/arrayobject.c
@@ -1762,6 +1762,59 @@ PyArray_SetMap(PyArrayMapIterObject *mit, PyObject *op)
 
 static PyObject *iter_subscript(PyArrayIterObject *, PyObject *); 
 
+int
+count_new_axes_0d(PyObject *tuple)
+{
+	int i, argument_count;
+	int ellipsis_count = 0;
+	int newaxis_count = 0;
+	argument_count = PyTuple_GET_SIZE(tuple);
+	for (i = 0; i < argument_count; ++i) {
+		PyObject *arg = PyTuple_GET_ITEM(tuple, i);
+		ellipsis_count += (arg == Py_Ellipsis);
+		newaxis_count += (arg == Py_None);
+	}
+	if (newaxis_count + ellipsis_count != argument_count) {
+		PyErr_SetString(PyExc_IndexError,
+				"0-d arrays can use a single ()"
+				" or a list of ellipses and newaxis"
+				" as an index");
+		return -1;
+	}
+	if (newaxis_count > MAX_DIMS) {
+		PyErr_SetString(PyExc_IndexError,
+				"too many dimensions");
+		return -1;
+	}
+	return newaxis_count;
+
+}
+static PyObject *
+add_new_axes_0d(PyArrayObject *arr,  int newaxis_count)
+{
+	PyArrayObject *other;
+	intp dimensions[MAX_DIMS], strides[MAX_DIMS];
+	int i;
+	for (i = 0; i < newaxis_count; ++i) {
+		dimensions[i] = strides[i] = 1;
+	}
+	Py_INCREF(arr->descr);
+	if ((other = (PyArrayObject *)
+	     PyArray_NewFromDescr(arr->ob_type, arr->descr,
+				  newaxis_count, dimensions,
+				  strides, arr->data,
+				  arr->flags,
+				  (PyObject *)arr)) == NULL) 
+		return NULL;
+
+	other->base = (PyObject *)arr;
+	Py_INCREF(arr);
+	
+	other ->flags &= ~OWNDATA;
+	
+	return (PyObject *)other;
+}
+
 static PyObject *
 array_subscript(PyArrayObject *self, PyObject *op) 
 {
@@ -1796,9 +1849,17 @@ array_subscript(PyArrayObject *self, PyObject *op)
 		return NULL;
 	}
         if (self->nd == 0) {
-		if (op == Py_Ellipsis || (PyTuple_Check(op) &&          \
-                                          0 == PyTuple_GET_SIZE(op)))
+		if (op == Py_Ellipsis)
 			return PyArray_ToScalar(self->data, self);
+		if (op == Py_None)
+			return add_new_axes_0d(self, 1);
+		if (PyTuple_Check(op)) {
+			if (0 == PyTuple_GET_SIZE(op))
+				return PyArray_ToScalar(self->data, self);
+			if ((nd = count_new_axes_0d(op)) == -1)
+				return NULL;
+			return add_new_axes_0d(self, nd);
+		}
                 PyErr_SetString(PyExc_IndexError, 
                                 "0-d arrays can't be indexed.");
                 return NULL;