Somewhat-preliminary slice-object and extended slicing support for ctypes.

The exact behaviour of omitted and negative indices for the Pointer type may
need a closer look (especially as it's subtly different from simple slices)
but there's time yet before 2.6, and not enough before 3.0a1 :-)

1 files changed, 305 insertions, 2 deletions

diff --git a/Modules/_ctypes/_ctypes.c b/Modules/_ctypes/_ctypes.c
index 020f473..4c7c896 100644
--- a/Modules/_ctypes/_ctypes.c
+++ b/Modules/_ctypes/_ctypes.c
@@ -3768,6 +3768,108 @@ Array_slice(PyObject *_self, Py_ssize_t ilow, Py_ssize_t ihigh)
 	return (PyObject *)np;
 }
 
+static PyObject *
+Array_subscript(PyObject *_self, PyObject *item)
+{
+	CDataObject *self = (CDataObject *)_self;
+
+	if (PyIndex_Check(item)) {
+		Py_ssize_t i = PyNumber_AsSsize_t(item, PyExc_IndexError);
+		
+		if (i == -1 && PyErr_Occurred())
+			return NULL;
+		if (i < 0)
+			i += self->b_length;
+		return Array_item(_self, i);
+	}
+	else if PySlice_Check(item) {
+		StgDictObject *stgdict, *itemdict;
+		PyObject *proto;
+		PyObject *np;
+		Py_ssize_t start, stop, step, slicelen, cur, i;
+		
+		if (PySlice_GetIndicesEx((PySliceObject *)item,
+					 self->b_length, &start, &stop,
+					 &step, &slicelen) < 0) {
+			return NULL;
+		}
+		
+		stgdict = PyObject_stgdict((PyObject *)self);
+		assert(stgdict); /* Cannot be NULL for array object instances */
+		proto = stgdict->proto;
+		itemdict = PyType_stgdict(proto);
+		assert(itemdict); /* proto is the item type of the array, a
+				     ctypes type, so this cannot be NULL */
+
+		if (itemdict->getfunc == getentry("c")->getfunc) {
+			char *ptr = (char *)self->b_ptr;
+			char *dest;
+
+			if (slicelen <= 0)
+				return PyString_FromString("");
+			if (step == 1) {
+				return PyString_FromStringAndSize(ptr + start,
+								  slicelen);
+			}
+			dest = (char *)PyMem_Malloc(slicelen);
+
+			if (dest == NULL)
+				return PyErr_NoMemory();
+
+			for (cur = start, i = 0; i < slicelen;
+			     cur += step, i++) {
+				dest[i] = ptr[cur];
+			}
+
+			np = PyString_FromStringAndSize(dest, slicelen);
+			PyMem_Free(dest);
+			return np;
+		}
+#ifdef CTYPES_UNICODE
+		if (itemdict->getfunc == getentry("u")->getfunc) {
+			wchar_t *ptr = (wchar_t *)self->b_ptr;
+			wchar_t *dest;
+			
+			if (slicelen <= 0)
+				return PyUnicode_FromUnicode(NULL, 0);
+			if (step == 1) {
+				return PyUnicode_FromWideChar(ptr + start,
+							      slicelen);
+			}
+
+			dest = (wchar_t *)PyMem_Malloc(
+						slicelen * sizeof(wchar_t));
+			
+			for (cur = start, i = 0; i < slicelen;
+			     cur += step, i++) {
+				dest[i] = ptr[cur];
+			}
+			
+			np = PyUnicode_FromWideChar(dest, slicelen);
+			PyMem_Free(dest);
+			return np;
+		}
+#endif
+
+		np = PyList_New(slicelen);
+		if (np == NULL)
+			return NULL;
+
+		for (cur = start, i = 0; i < slicelen;
+		     cur += step, i++) {
+			PyObject *v = Array_item(_self, cur);
+			PyList_SET_ITEM(np, i, v);
+		}
+		return np;
+	}
+	else {
+		PyErr_SetString(PyExc_TypeError, 
+				"indices must be integers");
+		return NULL;
+	}
+
+}
+
 static int
 Array_ass_item(PyObject *_self, Py_ssize_t index, PyObject *value)
 {
@@ -3839,6 +3941,63 @@ Array_ass_slice(PyObject *_self, Py_ssize_t ilow, Py_ssize_t ihigh, PyObject *va
 	return 0;
 }
 
+static int
+Array_ass_subscript(PyObject *_self, PyObject *item, PyObject *value)
+{
+	CDataObject *self = (CDataObject *)_self;
+	
+	if (value == NULL) {
+		PyErr_SetString(PyExc_TypeError,
+				"Array does not support item deletion");
+		return -1;
+	}
+
+	if (PyIndex_Check(item)) {
+		Py_ssize_t i = PyNumber_AsSsize_t(item, PyExc_IndexError);
+		
+		if (i == -1 && PyErr_Occurred())
+			return -1;
+		if (i < 0)
+			i += self->b_length;
+		return Array_ass_item(_self, i, value);
+	}
+	else if (PySlice_Check(item)) {
+		Py_ssize_t start, stop, step, slicelen, otherlen, i, cur;
+		
+		if (PySlice_GetIndicesEx((PySliceObject *)item,
+					 self->b_length, &start, &stop,
+					 &step, &slicelen) < 0) {
+			return -1;
+		}
+		if ((step < 0 && start < stop) ||
+		    (step > 0 && start > stop))
+			stop = start;
+
+		otherlen = PySequence_Length(value);
+		if (otherlen != slicelen) {
+			PyErr_SetString(PyExc_ValueError,
+				"Can only assign sequence of same size");
+			return -1;
+		}
+		for (cur = start, i = 0; i < otherlen; cur += step, i++) {
+			PyObject *item = PySequence_GetItem(value, i);
+			int result;
+			if (item == NULL)
+				return -1;
+			result = Array_ass_item(_self, cur, item);
+			Py_DECREF(item);
+			if (result == -1)
+				return -1;
+		}
+		return 0;
+	}
+	else {
+		PyErr_SetString(PyExc_TypeError,
+				"indices must be integer");
+		return -1;
+	}
+}
+
 static Py_ssize_t
 Array_length(PyObject *_self)
 {
@@ -3860,6 +4019,12 @@ static PySequenceMethods Array_as_sequence = {
 	0,					/* sq_inplace_repeat; */
 };
 
+static PyMappingMethods Array_as_mapping = {
+	Array_length,
+	Array_subscript,
+	Array_ass_subscript,
+};
+
 PyTypeObject Array_Type = {
 	PyVarObject_HEAD_INIT(NULL, 0)
 	"_ctypes.Array",
@@ -3873,7 +4038,7 @@ PyTypeObject Array_Type = {
 	0,					/* tp_repr */
 	0,					/* tp_as_number */
 	&Array_as_sequence,			/* tp_as_sequence */
-	0,					/* tp_as_mapping */
+	&Array_as_mapping,			/* tp_as_mapping */
 	0,					/* tp_hash */
 	0,					/* tp_call */
 	0,					/* tp_str */
@@ -4359,6 +4524,139 @@ Pointer_slice(PyObject *_self, Py_ssize_t ilow, Py_ssize_t ihigh)
 	return (PyObject *)np;
 }
 
+static PyObject *
+Pointer_subscript(PyObject *_self, PyObject *item)
+{
+	CDataObject *self = (CDataObject *)_self;
+	if (PyIndex_Check(item)) {
+		Py_ssize_t i = PyNumber_AsSsize_t(item, PyExc_IndexError);
+		if (i == -1 && PyErr_Occurred())
+			return NULL;
+		return Pointer_item(_self, i);
+	}
+	else if (PySlice_Check(item)) {
+		PySliceObject *slice = (PySliceObject *)item;
+		Py_ssize_t start, stop, step;
+		PyObject *np;
+		StgDictObject *stgdict, *itemdict;
+		PyObject *proto;
+		Py_ssize_t i, len, cur;
+
+		/* Since pointers have no length, and we want to apply
+		   different semantics to negative indices than normal
+		   slicing, we have to dissect the slice object ourselves.*/
+		if (slice->step == Py_None) {
+			step = 1;
+		}
+		else {
+			step = PyNumber_AsSsize_t(slice->step,
+						  PyExc_ValueError);
+			if (step == -1 && PyErr_Occurred())
+				return NULL;
+			if (step == 0) {
+				PyErr_SetString(PyExc_ValueError,
+						"slice step cannot be zero");
+				return NULL;
+			}
+		}
+		if (slice->start == Py_None) {
+			if (step < 0) {
+				PyErr_SetString(PyExc_ValueError,
+						"slice start is required "
+						"for step < 0");
+				return NULL;
+			}
+			start = 0;
+		}
+		else {
+			start = PyNumber_AsSsize_t(slice->start,
+						   PyExc_ValueError);
+			if (start == -1 && PyErr_Occurred())
+				return NULL;
+		}
+		if (slice->stop == Py_None) {
+			PyErr_SetString(PyExc_ValueError,
+					"slice stop is required");
+			return NULL;
+		}
+		stop = PyNumber_AsSsize_t(slice->stop,
+					  PyExc_ValueError);
+		if (stop == -1 && PyErr_Occurred())
+			return NULL;
+		if ((step > 0 && start > stop) ||
+		    (step < 0 && start < stop))
+			len = 0;
+		else if (step > 0)
+			len = (stop - start - 1) / step + 1;
+		else
+			len = (stop - start + 1) / step + 1;
+
+		stgdict = PyObject_stgdict((PyObject *)self);
+		assert(stgdict); /* Cannot be NULL for pointer instances */
+		proto = stgdict->proto;
+		assert(proto);
+		itemdict = PyType_stgdict(proto);
+		assert(itemdict);
+		if (itemdict->getfunc == getentry("c")->getfunc) {
+			char *ptr = *(char **)self->b_ptr;
+			char *dest;
+			
+			if (len <= 0)
+                        	return PyString_FromString("");
+			if (step == 1) {
+				return PyString_FromStringAndSize(ptr + start,
+								  len);
+			}
+			dest = (char *)PyMem_Malloc(len);
+			if (dest == NULL)
+				return PyErr_NoMemory();
+			for (cur = start, i = 0; i < len; cur += step, i++) {
+				dest[i] = ptr[cur];
+			}
+			np = PyString_FromStringAndSize(dest, len);
+			PyMem_Free(dest);
+			return np;
+		}
+#ifdef CTYPES_UNICODE
+		if (itemdict->getfunc == getentry("u")->getfunc) {
+			wchar_t *ptr = *(wchar_t **)self->b_ptr;
+			wchar_t *dest;
+			
+			if (len <= 0)
+                        	return PyUnicode_FromUnicode(NULL, 0);
+			if (step == 1) {
+				return PyUnicode_FromWideChar(ptr + start,
+							      len);
+			}
+			dest = (wchar_t *)PyMem_Malloc(len * sizeof(wchar_t));
+			if (dest == NULL)
+				return PyErr_NoMemory();
+			for (cur = start, i = 0; i < len; cur += step, i++) {
+				dest[i] = ptr[cur];
+			}
+			np = PyUnicode_FromWideChar(dest, len);
+			PyMem_Free(dest);
+			return np;
+		}
+#endif
+
+		np = PyList_New(len);
+		if (np == NULL)
+			return NULL;
+
+		for (cur = start, i = 0; i < len; cur += step, i++) {
+			PyObject *v = Pointer_item(_self, cur);
+			PyList_SET_ITEM(np, i, v);
+		}
+		return np;
+	}
+	else {
+		PyErr_SetString(PyExc_TypeError,
+				"Pointer indices must be integer");
+		return NULL;
+	}
+}
+
 static PySequenceMethods Pointer_as_sequence = {
 	0,					/* inquiry sq_length; */
 	0,					/* binaryfunc sq_concat; */
@@ -4373,6 +4671,11 @@ static PySequenceMethods Pointer_as_sequence = {
 	0,					/* intargfunc sq_inplace_repeat; */
 };
 
+static PyMappingMethods Pointer_as_mapping = {
+	0,
+	Pointer_subscript,
+};
+
 static int
 Pointer_nonzero(CDataObject *self)
 {
@@ -4406,7 +4709,7 @@ PyTypeObject Pointer_Type = {
 	0,					/* tp_repr */
 	&Pointer_as_number,			/* tp_as_number */
 	&Pointer_as_sequence,			/* tp_as_sequence */
-	0,					/* tp_as_mapping */
+	&Pointer_as_mapping,			/* tp_as_mapping */
 	0,					/* tp_hash */
 	0,					/* tp_call */
 	0,					/* tp_str */