#3156: fix consistency in what type bytearray methods accept as items.

Also rename confusing "item" parameters to "index".

3 files changed, 67 insertions, 56 deletions

diff --git a/Lib/test/test_bytes.py b/Lib/test/test_bytes.py
index d2d5b27..a33dcb5 100644
--- a/Lib/test/test_bytes.py
+++ b/Lib/test/test_bytes.py
@@ -17,6 +17,12 @@ import test.test_support
 import test.string_tests
 import test.buffer_tests
 
+class Indexable:
+    def __init__(self, value=0):
+        self.value = value
+    def __index__(self):
+        return self.value
+
 
 class BaseBytesTest(unittest.TestCase):
 
@@ -53,15 +59,11 @@ class BaseBytesTest(unittest.TestCase):
         self.assertEqual(list(b), ints)
 
     def test_from_index(self):
-        class C:
-            def __init__(self, i=0):
-                self.i = i
-            def __index__(self):
-                return self.i
-        b = self.type2test([C(), C(1), C(254), C(255)])
+        b = self.type2test([Indexable(), Indexable(1), Indexable(254),
+                            Indexable(255)])
         self.assertEqual(list(b), [0, 1, 254, 255])
-        self.assertRaises(ValueError, bytearray, [C(-1)])
-        self.assertRaises(ValueError, bytearray, [C(256)])
+        self.assertRaises(ValueError, bytearray, [Indexable(-1)])
+        self.assertRaises(ValueError, bytearray, [Indexable(256)])
 
     def test_from_ssize(self):
         self.assertEqual(bytearray(0), b'')
@@ -506,12 +508,7 @@ class ByteArrayTest(BaseBytesTest):
         self.assertEqual(b, bytearray([1, 100, 3]))
         b[-1] = 200
         self.assertEqual(b, bytearray([1, 100, 200]))
-        class C:
-            def __init__(self, i=0):
-                self.i = i
-            def __index__(self):
-                return self.i
-        b[0] = C(10)
+        b[0] = Indexable(10)
         self.assertEqual(b, bytearray([10, 100, 200]))
         try:
             b[3] = 0
@@ -529,7 +526,7 @@ class ByteArrayTest(BaseBytesTest):
         except ValueError:
             pass
         try:
-            b[0] = C(-1)
+            b[0] = Indexable(-1)
             self.fail("Didn't raise ValueError")
         except ValueError:
             pass
@@ -665,6 +662,9 @@ class ByteArrayTest(BaseBytesTest):
         self.assertRaises(ValueError, a.extend, [0, 1, 2, 256])
         self.assertRaises(ValueError, a.extend, [0, 1, 2, -1])
         self.assertEqual(len(a), 0)
+        a = bytearray(b'')
+        a.extend([Indexable(ord('a'))])
+        self.assertEqual(a, b'a')
 
     def test_remove(self):
         b = bytearray(b'hello')
@@ -680,6 +680,8 @@ class ByteArrayTest(BaseBytesTest):
         b.remove(ord('h'))
         self.assertEqual(b, b'e')
         self.assertRaises(TypeError, lambda: b.remove(u'e'))
+        b.remove(Indexable(ord('e')))
+        self.assertEqual(b, b'')
 
     def test_pop(self):
         b = bytearray(b'world')
@@ -701,6 +703,9 @@ class ByteArrayTest(BaseBytesTest):
         b.append(ord('A'))
         self.assertEqual(len(b), 1)
         self.assertRaises(TypeError, lambda: b.append(u'o'))
+        b = bytearray()
+        b.append(Indexable(ord('A')))
+        self.assertEqual(b, b'A')
 
     def test_insert(self):
         b = bytearray(b'msssspp')
@@ -709,7 +714,11 @@ class ByteArrayTest(BaseBytesTest):
         b.insert(-2, ord('i'))
         b.insert(1000, ord('i'))
         self.assertEqual(b, b'mississippi')
-        self.assertRaises(TypeError, lambda: b.insert(0, b'1'))
+        # allowed in 2.6
+        #self.assertRaises(TypeError, lambda: b.insert(0, b'1'))
+        b = bytearray()
+        b.insert(0, Indexable(ord('A')))
+        self.assertEqual(b, b'A')
 
     def test_partition_bytearray_doesnt_share_nullstring(self):
         a, b, c = bytearray(b"x").partition(b"y")
diff --git a/Misc/NEWS b/Misc/NEWS
index 1ee2b42..fdd4b89 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -10,6 +10,10 @@ What's New in Python 2.6 beta 2?
 Core and Builtins
 -----------------
 
+- Issue #3156: Fix inconsistent behavior of the bytearray type: all
+  its methods now allow for items objects that can be converted to
+  an integer using operator.index().
+
 - Issue #3360: Fix incorrect parsing of '020000000000.0', which
   produced a ValueError instead of giving the correct float.
 
diff --git a/Objects/bytearrayobject.c b/Objects/bytearrayobject.c
index 91a1ed9..b1f6962 100644
--- a/Objects/bytearrayobject.c
+++ b/Objects/bytearrayobject.c
@@ -34,22 +34,32 @@ _getbytevalue(PyObject* arg, int *value)
 {
     long face_value;
 
-    if (PyInt_Check(arg)) {
-        face_value = PyInt_AsLong(arg);
-        if (face_value < 0 || face_value >= 256) {
-            PyErr_SetString(PyExc_ValueError, "byte must be in range(0, 256)");
-            return 0;
-        }
-    }
-    else if (PyBytes_CheckExact(arg)) {
+    if (PyBytes_CheckExact(arg)) {
         if (Py_SIZE(arg) != 1) {
             PyErr_SetString(PyExc_ValueError, "string must be of size 1");
             return 0;
         }
-        face_value = Py_CHARMASK(((PyBytesObject*)arg)->ob_sval[0]);
+        *value = Py_CHARMASK(((PyBytesObject*)arg)->ob_sval[0]);
+        return 1;
+    }
+    else if (PyInt_Check(arg) || PyLong_Check(arg)) {
+        face_value = PyLong_AsLong(arg);
     }
     else {
-        PyErr_Format(PyExc_TypeError, "an integer or string of size 1 is required");
+        PyObject *index = PyNumber_Index(arg);
+        if (index == NULL) {
+            PyErr_Format(PyExc_TypeError,
+                         "an integer or string of size 1 is required");
+            return 0;
+        }
+        face_value = PyLong_AsLong(index);
+        Py_DECREF(index);
+    }
+    if (face_value == -1 && PyErr_Occurred())
+        return 0;
+
+    if (face_value < 0 || face_value >= 256) {
+        PyErr_SetString(PyExc_ValueError, "byte must be in range(0, 256)");
         return 0;
     }
 
@@ -410,10 +420,10 @@ bytes_getitem(PyByteArrayObject *self, Py_ssize_t i)
 }
 
 static PyObject *
-bytes_subscript(PyByteArrayObject *self, PyObject *item)
+bytes_subscript(PyByteArrayObject *self, PyObject *index)
 {
-    if (PyIndex_Check(item)) {
-        Py_ssize_t i = PyNumber_AsSsize_t(item, PyExc_IndexError);
+    if (PyIndex_Check(index)) {
+        Py_ssize_t i = PyNumber_AsSsize_t(index, PyExc_IndexError);
 
         if (i == -1 && PyErr_Occurred())
             return NULL;
@@ -427,9 +437,9 @@ bytes_subscript(PyByteArrayObject *self, PyObject *item)
         }
         return PyInt_FromLong((unsigned char)(self->ob_bytes[i]));
     }
-    else if (PySlice_Check(item)) {
+    else if (PySlice_Check(index)) {
         Py_ssize_t start, stop, step, slicelength, cur, i;
-        if (PySlice_GetIndicesEx((PySliceObject *)item,
+        if (PySlice_GetIndicesEx((PySliceObject *)index,
                                  PyByteArray_GET_SIZE(self),
                                  &start, &stop, &step, &slicelength) < 0) {
             return NULL;
@@ -574,13 +584,13 @@ bytes_setitem(PyByteArrayObject *self, Py_ssize_t i, PyObject *value)
 }
 
 static int
-bytes_ass_subscript(PyByteArrayObject *self, PyObject *item, PyObject *values)
+bytes_ass_subscript(PyByteArrayObject *self, PyObject *index, PyObject *values)
 {
     Py_ssize_t start, stop, step, slicelen, needed;
     char *bytes;
 
-    if (PyIndex_Check(item)) {
-        Py_ssize_t i = PyNumber_AsSsize_t(item, PyExc_IndexError);
+    if (PyIndex_Check(index)) {
+        Py_ssize_t i = PyNumber_AsSsize_t(index, PyExc_IndexError);
 
         if (i == -1 && PyErr_Occurred())
             return -1;
@@ -601,25 +611,15 @@ bytes_ass_subscript(PyByteArrayObject *self, PyObject *item, PyObject *values)
             slicelen = 1;
         }
         else {
-            Py_ssize_t ival = PyNumber_AsSsize_t(values, PyExc_ValueError);
-            if (ival == -1 && PyErr_Occurred()) {
-                int int_value;
-                /* Also accept str of size 1 in 2.x */
-                PyErr_Clear();
-                if (!_getbytevalue(values, &int_value))
-                    return -1;
-                ival = (int) int_value;
-            } else if (ival < 0 || ival >= 256) {
-                PyErr_SetString(PyExc_ValueError,
-                                "byte must be in range(0, 256)");
+            int ival;
+            if (!_getbytevalue(values, &ival))
                 return -1;
-            }
             self->ob_bytes[i] = (char)ival;
             return 0;
         }
     }
-    else if (PySlice_Check(item)) {
-        if (PySlice_GetIndicesEx((PySliceObject *)item,
+    else if (PySlice_Check(index)) {
+        if (PySlice_GetIndicesEx((PySliceObject *)index,
                                  PyByteArray_GET_SIZE(self),
                                  &start, &stop, &step, &slicelen) < 0) {
             return -1;
@@ -640,7 +640,7 @@ bytes_ass_subscript(PyByteArrayObject *self, PyObject *item, PyObject *values)
         values = PyByteArray_FromObject(values);
         if (values == NULL)
             return -1;
-        err = bytes_ass_subscript(self, item, values);
+        err = bytes_ass_subscript(self, index, values);
         Py_DECREF(values);
         return err;
     }
@@ -2594,10 +2594,11 @@ Insert a single item into the bytearray before the given index.");
 static PyObject *
 bytes_insert(PyByteArrayObject *self, PyObject *args)
 {
-    int value;
+    PyObject *value;
+    int ival;
     Py_ssize_t where, n = Py_SIZE(self);
 
-    if (!PyArg_ParseTuple(args, "ni:insert", &where, &value))
+    if (!PyArg_ParseTuple(args, "nO:insert", &where, &value))
         return NULL;
 
     if (n == PY_SSIZE_T_MAX) {
@@ -2605,11 +2606,8 @@ bytes_insert(PyByteArrayObject *self, PyObject *args)
                         "cannot add more objects to bytes");
         return NULL;
     }
-    if (value < 0 || value >= 256) {
-        PyErr_SetString(PyExc_ValueError,
-                        "byte must be in range(0, 256)");
+    if (!_getbytevalue(value, &ival))
         return NULL;
-    }
     if (PyByteArray_Resize((PyObject *)self, n + 1) < 0)
         return NULL;
 
@@ -2621,7 +2619,7 @@ bytes_insert(PyByteArrayObject *self, PyObject *args)
     if (where > n)
         where = n;
     memmove(self->ob_bytes + where + 1, self->ob_bytes + where, n - where);
-    self->ob_bytes[where] = value;
+    self->ob_bytes[where] = ival;
 
     Py_RETURN_NONE;
 }