SF feature request #686323: Minor array module enhancements

rhettinger · rhettinger · commit 49f9bd15ffee · 2004-03-14T05:43:59.000Z
array.extend() now accepts iterable arguments implements as a series
of appends.  Besides being a user convenience and matching the behavior
for lists, this the saves memory and cycles that would be used to
create a temporary array object.
diff --git a/Doc/lib/libarray.tex b/Doc/lib/libarray.tex
@@ -104,10 +104,13 @@ \section{\module{array} ---
 Return the number of occurences of \var{x} in the array.
 \end{methoddesc}
 
-\begin{methoddesc}[array]{extend}{a}
-Append array items from \var{a} to the end of the array.  The two
-arrays must have \emph{exactly} the same type code; if not,
-\exception{TypeError} will be raised.
+\begin{methoddesc}[array]{extend}{iterable}
+Append items from \var{iterable} to the end of the array.  If
+\var{iterable} is another array, it must have \emph{exactly} the same
+type code; if not, \exception{TypeError} will be raised.  If
+\var{iterable} is not an array, it must be iterable and its
+elements must be the right type to be appended to the array.
+\versionchanged[Formerly, the argument could only be another array]{2.4}
 \end{methoddesc}
 
 \begin{methoddesc}[array]{fromfile}{f, n}
diff --git a/Lib/test/test_array.py b/Lib/test/test_array.py
@@ -592,6 +592,13 @@ def test_extend(self):
         b = array.array(self.badtypecode())
         self.assertRaises(TypeError, a.extend, b)
 
+        a = array.array(self.typecode, self.example)
+        a.extend(self.example[::-1])
+        self.assertEqual(
+            a,
+            array.array(self.typecode, self.example+self.example[::-1])
+        )
+
     def test_coveritertraverse(self):
         try:
             import gc
diff --git a/Misc/NEWS b/Misc/NEWS
@@ -182,7 +182,9 @@ Extension modules
 
 - array objects now support the copy module.  Also, their resizing
   scheme has been updated the same as for list objects.  The improves
-  performance for append() operations.
+  the performance (speed and memory usage) of append() operations.
+  Also, array.extend() now accepts any iterable argument for repeated
+  appends without needing to create another temporary array.
 
 - cStringIO.writelines() now accepts any iterable argument and writes
   the lines one at a time rather than joining them and writing once.
diff --git a/Modules/arraymodule.c b/Modules/arraymodule.c
@@ -775,16 +775,35 @@ setarrayitem(PyObject *a, int i, PyObject *v)
 }
 
 static int
-array_do_extend(arrayobject *self, PyObject *bb)
+array_iter_extend(arrayobject *self, PyObject *bb)
 {
-	int size;
+	PyObject *it, *v;
 
-	if (!array_Check(bb)) {
-		PyErr_Format(PyExc_TypeError,
-			"can only extend array with array (not \"%.200s\")",
-			bb->ob_type->tp_name);
+	it = PyObject_GetIter(bb);
+	if (it == NULL)
 		return -1;
+
+	while ((v = PyIter_Next(it)) != NULL) {
+		if (ins1(self, (int) self->ob_size, v) != 0) {
+			Py_DECREF(v);
+			Py_DECREF(it);
+			return -1;
+		}
+		Py_DECREF(v);
 	}
+	Py_DECREF(it);
+	if (PyErr_Occurred())
+		return -1;
+	return 0;
+}
+
+static int
+array_do_extend(arrayobject *self, PyObject *bb)
+{
+	int size;
+
+	if (!array_Check(bb))
+		return array_iter_extend(self, bb);
 #define b ((arrayobject *)bb)
 	if (self->ob_descr != b->ob_descr) {
 		PyErr_SetString(PyExc_TypeError,
@@ -810,6 +829,12 @@ array_do_extend(arrayobject *self, PyObject *bb)
 static PyObject *
 array_inplace_concat(arrayobject *self, PyObject *bb)
 {
+	if (!array_Check(bb)) {
+		PyErr_Format(PyExc_TypeError,
+			"can only extend array with array (not \"%.200s\")",
+			bb->ob_type->tp_name);
+		return NULL;
+	}
 	if (array_do_extend(self, bb) == -1)
 		return NULL;
 	Py_INCREF(self);
@@ -990,9 +1015,9 @@ array_extend(arrayobject *self, PyObject *bb)
 }
 
 PyDoc_STRVAR(extend_doc,
-"extend(array)\n\
+"extend(array or iterable)\n\
 \n\
- Append array items to the end of the array.");
+ Append items to the end of the array.");
 
 static PyObject *
 array_insert(arrayobject *self, PyObject *args)
@@ -1881,7 +1906,7 @@ append() -- append a new item to the end of the array\n\
 buffer_info() -- return information giving the current memory info\n\
 byteswap() -- byteswap all the items of the array\n\
 count() -- return number of occurences of an object\n\
-extend() -- extend array by appending array elements\n\
+extend() -- extend array by appending multiple elements from an iterable\n\
 fromfile() -- read items from a file object\n\
 fromlist() -- append items from the list\n\
 fromstring() -- append items from the string\n\
