Issue #1283: Allow any iterable of integers to be passed to

bytearray.extend().

2 files changed, 89 insertions, 18 deletions

diff --git a/Lib/test/test_bytes.py b/Lib/test/test_bytes.py
index 659afac..f037d4c 100644
--- a/Lib/test/test_bytes.py
+++ b/Lib/test/test_bytes.py
@@ -529,6 +529,24 @@ class BytesTest(unittest.TestCase):
         a.extend(a)
         self.assertEqual(a, orig + orig)
         self.assertEqual(a[5:], orig)
+        a = bytearray(b'')
+        # Test iterators that don't have a __length_hint__
+        a.extend(map(int, orig * 25))
+        a.extend(int(x) for x in orig * 25)
+        self.assertEqual(a, orig * 50)
+        self.assertEqual(a[-5:], orig)
+        a = bytearray(b'')
+        a.extend(iter(map(int, orig * 50)))
+        self.assertEqual(a, orig * 50)
+        self.assertEqual(a[-5:], orig)
+        a = bytearray(b'')
+        a.extend(list(map(int, orig * 50)))
+        self.assertEqual(a, orig * 50)
+        self.assertEqual(a[-5:], orig)
+        a = bytearray(b'')
+        self.assertRaises(ValueError, a.extend, [0, 1, 2, 256])
+        self.assertRaises(ValueError, a.extend, [0, 1, 2, -1])
+        self.assertEqual(len(a), 0)
 
     def test_remove(self):
         b = bytearray(b'hello')
diff --git a/Objects/bytesobject.c b/Objects/bytesobject.c
index fae29e2..021925a 100644
--- a/Objects/bytesobject.c
+++ b/Objects/bytesobject.c
@@ -2487,24 +2487,6 @@ onError:
     return NULL;
 }
 
-PyDoc_STRVAR(extend__doc__,
-"B.extend(iterable int) -> None\n\
-\n\
-Append all the elements from the iterator or sequence to the\n\
-end of B.");
-static PyObject *
-bytes_extend(PyBytesObject *self, PyObject *arg)
-{
-    /* XXX(gps): The docstring says any iterable int will do but the
-     * bytes_setslice code only accepts something supporting PEP 3118.
-     * A list or tuple of 0 <= int <= 255 is supposed to work.  */
-    /* bug being tracked on:  http://bugs.python.org/issue1283  */
-    if (bytes_setslice(self, Py_Size(self), Py_Size(self), arg) == -1)
-        return NULL;
-    Py_RETURN_NONE;
-}
-
-
 PyDoc_STRVAR(reverse__doc__,
 "B.reverse() -> None\n\
 \n\
@@ -2591,6 +2573,77 @@ bytes_append(PyBytesObject *self, PyObject *arg)
     Py_RETURN_NONE;
 }
 
+PyDoc_STRVAR(extend__doc__,
+"B.extend(iterable int) -> None\n\
+\n\
+Append all the elements from the iterator or sequence to the\n\
+end of B.");
+static PyObject *
+bytes_extend(PyBytesObject *self, PyObject *arg)
+{
+    PyObject *it, *item, *tmp, *res;
+    Py_ssize_t buf_size = 0, len = 0;
+    int value;
+    char *buf;
+
+    /* bytes_setslice code only accepts something supporting PEP 3118. */
+    if (PyObject_CheckBuffer(arg)) {
+        if (bytes_setslice(self, Py_Size(self), Py_Size(self), arg) == -1)
+            return NULL;
+
+        Py_RETURN_NONE;
+    }
+
+    it = PyObject_GetIter(arg);
+    if (it == NULL)
+        return NULL;
+
+    /* Try to determine the length of the argument. */
+    buf_size = _PyObject_LengthHint(arg);
+    /* The length of the argument is unknown or invalid. */
+    if (buf_size < 0) {
+		if (PyErr_Occurred()
+            && !PyErr_ExceptionMatches(PyExc_TypeError)
+            && !PyErr_ExceptionMatches(PyExc_AttributeError)) {
+            Py_DECREF(it);
+			return NULL;
+		}
+        PyErr_Clear();
+        buf_size = 32;          /* arbitrary */
+    }
+
+    buf = (char *)PyMem_Malloc(buf_size * sizeof(char));
+    if (buf == NULL)
+        return PyErr_NoMemory();
+
+    while ((item = PyIter_Next(it)) != NULL) {
+        if (! _getbytevalue(item, &value)) {
+            Py_DECREF(item);
+            Py_DECREF(it);
+            return NULL;
+        }
+        buf[len++] = value;
+        Py_DECREF(item);
+        if (len >= buf_size) {
+            buf_size = len + (len >> 1) + 1;
+            buf = (char *)PyMem_Realloc(buf, buf_size * sizeof(char));
+            if (buf == NULL) {
+                Py_DECREF(it);
+                return PyErr_NoMemory();
+            }
+        }
+    }
+    Py_DECREF(it);
+
+    /* XXX: Is possible to avoid a full copy of the buffer? */
+    tmp = PyBytes_FromStringAndSize(buf, len);
+    res = bytes_extend(self, tmp);
+    Py_DECREF(tmp);
+    PyMem_Free(buf);
+
+    return res;
+}
+
 PyDoc_STRVAR(pop__doc__,
 "B.pop([index]) -> int\n\
 \n\