Improve extended slicing support in builtin types and classes. Specifically:

 - Specialcase extended slices that amount to a shallow copy the same way as
   is done for simple slices, in the tuple, string and unicode case.

 - Specialcase step-1 extended slices to optimize the common case for all
   involved types.

 - For lists, allow extended slice assignment of differing lengths as long
   as the step is 1. (Previously, 'l[:2:1] = []' failed even though
   'l[:2] = []' and 'l[:2:None] = []' do not.)

 - Implement extended slicing for buffer, array, structseq, mmap and
   UserString.UserString.

 - Implement slice-object support (but not non-step-1 slice assignment) for
   UserString.MutableString.

 - Add tests for all new functionality.

2 files changed, 287 insertions, 93 deletions

diff --git a/Modules/arraymodule.c b/Modules/arraymodule.c
index 21a5e5b..b2ee5b0 100644
--- a/Modules/arraymodule.c
+++ b/Modules/arraymodule.c
@@ -1605,6 +1605,16 @@ array_subscr(arrayobject* self, PyObject* item)
 		if (slicelength <= 0) {
 			return newarrayobject(&Arraytype, 0, self->ob_descr);
 		}
+		else if (step == 1) {
+			PyObject *result = newarrayobject(&Arraytype,
+						slicelength, self->ob_descr);
+			if (result == NULL)
+				return NULL;
+			memcpy(((arrayobject *)result)->ob_item,
+			       self->ob_item + start * itemsize,
+			       slicelength * itemsize);
+			return result;
+		}
 		else {
 			result = newarrayobject(&Arraytype, slicelength, self->ob_descr);
 			if (!result) return NULL;
@@ -1623,7 +1633,7 @@ array_subscr(arrayobject* self, PyObject* item)
 	}
 	else {
 		PyErr_SetString(PyExc_TypeError, 
-				"list indices must be integers");
+				"array indices must be integers");
 		return NULL;
 	}
 }
@@ -1631,112 +1641,146 @@ array_subscr(arrayobject* self, PyObject* item)
 static int
 array_ass_subscr(arrayobject* self, PyObject* item, PyObject* value)
 {
+	Py_ssize_t start, stop, step, slicelength, needed;
+	arrayobject* other;
+	int itemsize;
+
 	if (PyIndex_Check(item)) {
 		Py_ssize_t i = PyNumber_AsSsize_t(item, PyExc_IndexError);
-		if (i==-1 && PyErr_Occurred()) 
+		
+		if (i == -1 && PyErr_Occurred())
 			return -1;
 		if (i < 0)
 			i += Py_Size(self);
-		return array_ass_item(self, i, value);
-	}
-	else if (PySlice_Check(item)) {
-		Py_ssize_t start, stop, step, slicelength;
-		int itemsize = self->ob_descr->itemsize;
-
-		if (PySlice_GetIndicesEx((PySliceObject*)item, Py_Size(self),
-				 &start, &stop, &step, &slicelength) < 0) {
+		if (i < 0 || i >= Py_Size(self)) {
+			PyErr_SetString(PyExc_IndexError,
+				"array assignment index out of range");
 			return -1;
 		}
-
-		/* treat A[slice(a,b)] = v _exactly_ like A[a:b] = v */
-		if (step == 1 && ((PySliceObject*)item)->step == Py_None)
-			return array_ass_slice(self, start, stop, value);
-
 		if (value == NULL) {
-			/* delete slice */
-			Py_ssize_t cur, i, extra;
-			
-			if (slicelength <= 0)
-				return 0;
-
-			if (step < 0) {
-				stop = start + 1;
-				start = stop + step*(slicelength - 1) - 1;
-				step = -step;
-			}
-
-			for (cur = start, i = 0; i < slicelength - 1;
-			     cur += step, i++) {
-				memmove(self->ob_item + (cur - i)*itemsize,
-					self->ob_item + (cur + 1)*itemsize,
-					(step - 1) * itemsize);
-			}
-			extra = Py_Size(self) - (cur + 1);
-			if (extra > 0) {
-				memmove(self->ob_item + (cur - i)*itemsize,
-					self->ob_item + (cur + 1)*itemsize,
-					extra*itemsize);
-			}
-
-			Py_Size(self) -= slicelength;
-			self->ob_item = (char *)PyMem_REALLOC(self->ob_item,
-							      itemsize*Py_Size(self));
-			self->allocated = Py_Size(self);
-
-			return 0;
+			/* Fall through to slice assignment */
+			start = i;
+			stop = i + 1;
+			step = 1;
+			slicelength = 1;
 		}
-		else {
-			/* assign slice */
-			Py_ssize_t cur, i;
-			arrayobject* av;
-
-			if (!array_Check(value)) {
-				PyErr_Format(PyExc_TypeError,
-			     "must assign array (not \"%.200s\") to slice",
-					     Py_Type(value)->tp_name);
-				return -1;
-			}
-
-			av = (arrayobject*)value;
-
-			if (Py_Size(av) != slicelength) {
-				PyErr_Format(PyExc_ValueError,
-            "attempt to assign array of size %ld to extended slice of size %ld",
-					     /*XXX*/(long)Py_Size(av), /*XXX*/(long)slicelength);
+		else
+			return (*self->ob_descr->setitem)(self, i, value);
+	}
+	else if (PySlice_Check(item)) {
+		if (PySlice_GetIndicesEx((PySliceObject *)item,
+					 Py_Size(self), &start, &stop,
+					 &step, &slicelength) < 0) {
+			return -1;
+		}
+	}
+	else {
+		PyErr_SetString(PyExc_TypeError,
+				"array indices must be integer");
+		return -1;
+	}
+	if (value == NULL) {
+		other = NULL;
+		needed = 0;
+	}
+	else if (array_Check(value)) {
+		other = (arrayobject *)value;
+		needed = Py_Size(other);
+		if (self == other) {
+			/* Special case "self[i:j] = self" -- copy self first */
+			int ret;
+			value = array_slice(other, 0, needed);
+			if (value == NULL)
 				return -1;
-			}
-
-			if (!slicelength)
-				return 0;
-
-			/* protect against a[::-1] = a */
-			if (self == av) { 
-				value = array_slice(av, 0, Py_Size(av));
-				av = (arrayobject*)value;
-				if (!av)
-					return -1;
-			} 
-			else {
-				Py_INCREF(value);
-			}
-
-			for (cur = start, i = 0; i < slicelength; 
-			     cur += step, i++) {
-				memcpy(self->ob_item + cur*itemsize,
-				       av->ob_item + i*itemsize,
-				       itemsize);
-			}
-
+			ret = array_ass_subscr(self, item, value);
 			Py_DECREF(value);
-			
-			return 0;
+			return ret;
+		}
+		if (other->ob_descr != self->ob_descr) {
+			PyErr_BadArgument();
+			return -1;
 		}
-	} 
+	}
 	else {
-		PyErr_SetString(PyExc_TypeError, 
-				"list indices must be integers");
+		PyErr_Format(PyExc_TypeError,
+	     "can only assign array (not \"%.200s\") to array slice",
+			     Py_Type(value)->tp_name);
 		return -1;
 	}
+	itemsize = self->ob_descr->itemsize;
+	/* for 'a[2:1] = ...', the insertion point is 'start', not 'stop' */
+	if ((step > 0 && stop < start) ||
+	    (step < 0 && stop > start))
+		stop = start;
+	if (step == 1) {
+		if (slicelength > needed) {
+			memmove(self->ob_item + (start + needed) * itemsize,
+				self->ob_item + stop * itemsize,
+				(Py_Size(self) - stop) * itemsize);
+			if (array_resize(self, Py_Size(self) +
+					 needed - slicelength) < 0)
+				return -1;
+		}
+		else if (slicelength < needed) {
+			if (array_resize(self, Py_Size(self) +
+					 needed - slicelength) < 0)
+				return -1;
+			memmove(self->ob_item + (start + needed) * itemsize,
+				self->ob_item + stop * itemsize,
+				(Py_Size(self) - start - needed) * itemsize);
+		}
+		if (needed > 0)
+			memcpy(self->ob_item + start * itemsize,
+			       other->ob_item, needed * itemsize);
+		return 0;
+	}
+	else if (needed == 0) {
+		/* Delete slice */
+		Py_ssize_t cur, i;
+		
+		if (step < 0) {
+			stop = start + 1;
+			start = stop + step * (slicelength - 1) - 1;
+			step = -step;
+		}
+		for (cur = start, i = 0; i < slicelength;
+		     cur += step, i++) {
+			Py_ssize_t lim = step - 1;
+
+			if (cur + step >= Py_Size(self))
+				lim = Py_Size(self) - cur - 1;
+			memmove(self->ob_item + (cur - i) * itemsize,
+				self->ob_item + (cur + 1) * itemsize,
+				lim * itemsize);
+		}
+		cur = start + slicelength * step;
+		if (cur < Py_Size(self)) {
+			memmove(self->ob_item + (cur-slicelength) * itemsize,
+				self->ob_item + cur * itemsize,
+				(Py_Size(self) - cur) * itemsize);
+		}
+		if (array_resize(self, Py_Size(self) - slicelength) < 0)
+			return -1;
+		return 0;
+	}
+	else {
+		Py_ssize_t cur, i;
+
+		if (needed != slicelength) {
+			PyErr_Format(PyExc_ValueError,
+				"attempt to assign array of size %zd "
+				"to extended slice of size %zd",
+				needed, slicelength);
+			return -1;
+		}
+		for (cur = start, i = 0; i < slicelength;
+		     cur += step, i++) {
+			memcpy(self->ob_item + cur * itemsize,
+			       other->ob_item + i * itemsize,
+			       itemsize);
+		}
+		return 0;
+	}
 }
 
 static PyMappingMethods array_as_mapping = {
diff --git a/Modules/mmapmodule.c b/Modules/mmapmodule.c
index a92954b..6114442 100644
--- a/Modules/mmapmodule.c
+++ b/Modules/mmapmodule.c
@@ -681,6 +681,60 @@ mmap_slice(mmap_object *self, Py_ssize_t ilow, Py_ssize_t ihigh)
 }
 
 static PyObject *
+mmap_subscript(mmap_object *self, PyObject *item)
+{
+	CHECK_VALID(NULL);
+	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->size;
+		if (i < 0 || i > self->size) {
+			PyErr_SetString(PyExc_IndexError,
+				"mmap index out of range");
+			return NULL;
+		}
+		return PyString_FromStringAndSize(self->data + i, 1);
+	}
+	else if (PySlice_Check(item)) {
+		Py_ssize_t start, stop, step, slicelen;
+
+		if (PySlice_GetIndicesEx((PySliceObject *)item, self->size,
+				 &start, &stop, &step, &slicelen) < 0) {
+			return NULL;
+		}
+		
+		if (slicelen <= 0)
+			return PyString_FromStringAndSize("", 0);
+		else if (step == 1)
+			return PyString_FromStringAndSize(self->data + start,
+							  slicelen);
+		else {
+			char *result_buf = (char *)PyMem_Malloc(slicelen);
+			Py_ssize_t cur, i;
+			PyObject *result;
+
+			if (result_buf == NULL)
+				return PyErr_NoMemory();
+			for (cur = start, i = 0; i < slicelen;
+			     cur += step, i++) {
+			     	result_buf[i] = self->data[cur];
+			}
+			result = PyString_FromStringAndSize(result_buf,
+							    slicelen);
+			PyMem_Free(result_buf);
+			return result;
+		}
+	}
+	else {
+		PyErr_SetString(PyExc_TypeError,
+				"mmap indices must be integers");
+		return NULL;
+	}
+}
+
+static PyObject *
 mmap_concat(mmap_object *self, PyObject *bb)
 {
 	CHECK_VALID(NULL);
@@ -764,6 +818,96 @@ mmap_ass_item(mmap_object *self, Py_ssize_t i, PyObject *v)
 	return 0;
 }
 
+static int
+mmap_ass_subscript(mmap_object *self, PyObject *item, PyObject *value)
+{
+	CHECK_VALID(-1);
+
+	if (PyIndex_Check(item)) {
+		Py_ssize_t i = PyNumber_AsSsize_t(item, PyExc_IndexError);
+		const char *buf;
+
+		if (i == -1 && PyErr_Occurred())
+			return -1;
+		if (i < 0)
+			i += self->size;
+		if (i < 0 || i > self->size) {
+			PyErr_SetString(PyExc_IndexError,
+				"mmap index out of range");
+			return -1;
+		}
+		if (value == NULL) {
+			PyErr_SetString(PyExc_TypeError,
+				"mmap object doesn't support item deletion");
+			return -1;
+		}
+		if (!PyString_Check(value) || PyString_Size(value) != 1) {
+			PyErr_SetString(PyExc_IndexError,
+		          "mmap assignment must be single-character string");
+			return -1;
+		}
+		if (!is_writeable(self))
+			return -1;
+		buf = PyString_AsString(value);
+		self->data[i] = buf[0];
+		return 0;
+	}
+	else if (PySlice_Check(item)) {
+		Py_ssize_t start, stop, step, slicelen;
+		
+		if (PySlice_GetIndicesEx((PySliceObject *)item,
+					 self->size, &start, &stop,
+					 &step, &slicelen) < 0) {
+			return -1;
+		}
+		if (value == NULL) {
+			PyErr_SetString(PyExc_TypeError,
+				"mmap object doesn't support slice deletion");
+			return -1;
+		}
+		if (!PyString_Check(value)) {
+			PyErr_SetString(PyExc_IndexError,
+				"mmap slice assignment must be a string");
+			return -1;
+		}
+		if (PyString_Size(value) != slicelen) {
+			PyErr_SetString(PyExc_IndexError,
+				"mmap slice assignment is wrong size");
+			return -1;
+		}
+		if (!is_writeable(self))
+			return -1;
+
+		if (slicelen == 0)
+			return 0;
+		else if (step == 1) {
+			const char *buf = PyString_AsString(value);
+
+			if (buf == NULL)
+				return -1;
+			memcpy(self->data + start, buf, slicelen);
+			return 0;
+		}
+		else {
+			Py_ssize_t cur, i;
+			const char *buf = PyString_AsString(value);
+			
+			if (buf == NULL)
+				return -1;
+			for (cur = start, i = 0; i < slicelen;
+			     cur += step, i++) {
+				self->data[cur] = buf[i];
+			}
+			return 0;
+		}
+	}
+	else {
+		PyErr_SetString(PyExc_TypeError,
+				"mmap indices must be integer");
+		return -1;
+	}
+}
+
 static PySequenceMethods mmap_as_sequence = {
 	(lenfunc)mmap_length,		       /*sq_length*/
 	(binaryfunc)mmap_concat,	       /*sq_concat*/
@@ -774,6 +918,12 @@ static PySequenceMethods mmap_as_sequence = {
 	(ssizessizeobjargproc)mmap_ass_slice,      /*sq_ass_slice*/
 };
 
+static PyMappingMethods mmap_as_mapping = {
+	(lenfunc)mmap_length,
+	(binaryfunc)mmap_subscript,
+	(objobjargproc)mmap_ass_subscript,
+};
+
 static PyBufferProcs mmap_as_buffer = {
 	(readbufferproc)mmap_buffer_getreadbuf,
 	(writebufferproc)mmap_buffer_getwritebuf,
@@ -795,7 +945,7 @@ static PyTypeObject mmap_object_type = {
 	0,					/* tp_repr */
 	0,					/* tp_as_number */
 	&mmap_as_sequence,			/*tp_as_sequence*/
-	0,					/*tp_as_mapping*/
+	&mmap_as_mapping,			/*tp_as_mapping*/
 	0,					/*tp_hash*/
 	0,					/*tp_call*/
 	0,					/*tp_str*/