bpo-28298: make array 'Q', 'L' and 'I' accept big intables as elements

Author: orenmn

Lib/test/test_array.py

@@ -1225,7 +1225,26 @@ def test_frombytearray(self):
         b = array.array(self.typecode, a)
         self.assertEqual(a, b)
 
-class SignedNumberTest(NumberTest):
+class IntegerNumberTest(NumberTest):
+    def test_type_error(self):
+        a = array.array(self.typecode)
+        a.append(42)
+        with self.assertRaises(TypeError):
+            a.append(42.0)
+        with self.assertRaises(TypeError):
+            a[0] = 42.0
+
+class Intable:
+    def __init__(self, num):
+        self._num = num
+    def __int__(self):
+        return self._num
+    def __sub__(self, other):
+        return Intable(int(self) - int(other))
+    def __add__(self, other):
+        return Intable(int(self) + int(other))
+
+class SignedNumberTest(IntegerNumberTest):
     example = [-1, 0, 1, 42, 0x7f]
     smallerexample = [-1, 0, 1, 42, 0x7e]
     biggerexample = [-1, 0, 1, 43, 0x7f]
@@ -1236,8 +1255,9 @@ def test_overflow(self):
         lower = -1 * int(pow(2, a.itemsize * 8 - 1))
         upper = int(pow(2, a.itemsize * 8 - 1)) - 1
         self.check_overflow(lower, upper)
+        self.check_overflow(Intable(lower), Intable(upper))
 
-class UnsignedNumberTest(NumberTest):
+class UnsignedNumberTest(IntegerNumberTest):
     example = [0, 1, 17, 23, 42, 0xff]
     smallerexample = [0, 1, 17, 23, 42, 0xfe]
     biggerexample = [0, 1, 17, 23, 43, 0xff]
@@ -1248,6 +1268,7 @@ def test_overflow(self):
         lower = 0
         upper = int(pow(2, a.itemsize * 8)) - 1
         self.check_overflow(lower, upper)
+        self.check_overflow(Intable(lower), Intable(upper))
 
     def test_bytes_extend(self):
         s = bytes(self.example)

Misc/NEWS

@@ -41,6 +41,9 @@ Extension Modules
 Library
 -------
 
+- bpo-28298: Fix a bug that prevented array 'Q', 'L' and 'I' from accepting big
+  intables (objects that have __int__) as elements.  Patch by Oren Milman.
+
 - bpo-29615: SimpleXMLRPCDispatcher no longer chains KeyError (or any other
   exception) to exception(s) raised in the dispatched methods.
   Patch by Petr Motejlek.

Modules/arraymodule.c

@@ -331,35 +331,51 @@ II_getitem(arrayobject *ap, Py_ssize_t i)
         (unsigned long) ((unsigned int *)ap->ob_item)[i]);
 }
 
+static PyObject *
+get_int_unless_float(PyObject *v)
+{
+    if (PyFloat_Check(v)) {
+        PyErr_SetString(PyExc_TypeError,
+                        "array item must be integer");
+        return NULL;
+    }
+    return (PyObject *)_PyLong_FromNbInt(v);
+}
+
 static int
 II_setitem(arrayobject *ap, Py_ssize_t i, PyObject *v)
 {
     unsigned long x;
-    if (PyLong_Check(v)) {
-        x = PyLong_AsUnsignedLong(v);
-        if (x == (unsigned long) -1 && PyErr_Occurred())
+    int do_decref = 0; /* if nb_int was called */
+
+    if (!PyLong_Check(v)) {
+        v = get_int_unless_float(v);
+        if (NULL == v) {
             return -1;
+        }
+        do_decref = 1;
     }
-    else {
-        long y;
-        if (!PyArg_Parse(v, "l;array item must be integer", &y))
-            return -1;
-        if (y < 0) {
-            PyErr_SetString(PyExc_OverflowError,
-                "unsigned int is less than minimum");
-            return -1;
+    x = PyLong_AsUnsignedLong(v);
+    if (x == (unsigned long)-1 && PyErr_Occurred()) {
+        if (do_decref) {
+            Py_DECREF(v);
         }
-        x = (unsigned long)y;
-
+        return -1;
     }
     if (x > UINT_MAX) {
         PyErr_SetString(PyExc_OverflowError,
-            "unsigned int is greater than maximum");
+                        "unsigned int is greater than maximum");
+        if (do_decref) {
+            Py_DECREF(v);
+        }
         return -1;
     }
-
     if (i >= 0)
         ((unsigned int *)ap->ob_item)[i] = (unsigned int)x;
+
+    if (do_decref) {
+        Py_DECREF(v);
+    }
     return 0;
 }
 
@@ -390,31 +406,28 @@ static int
 LL_setitem(arrayobject *ap, Py_ssize_t i, PyObject *v)
 {
     unsigned long x;
-    if (PyLong_Check(v)) {
-        x = PyLong_AsUnsignedLong(v);
-        if (x == (unsigned long) -1 && PyErr_Occurred())
-            return -1;
-    }
-    else {
-        long y;
-        if (!PyArg_Parse(v, "l;array item must be integer", &y))
-            return -1;
-        if (y < 0) {
-            PyErr_SetString(PyExc_OverflowError,
-                "unsigned long is less than minimum");
+    int do_decref = 0; /* if nb_int was called */
+
+    if (!PyLong_Check(v)) {
+        v = get_int_unless_float(v);
+        if (NULL == v) {
             return -1;
         }
-        x = (unsigned long)y;
-
+        do_decref = 1;
     }
-    if (x > ULONG_MAX) {
-        PyErr_SetString(PyExc_OverflowError,
-            "unsigned long is greater than maximum");
+    x = PyLong_AsUnsignedLong(v);
+    if (x == (unsigned long)-1 && PyErr_Occurred()) {
+        if (do_decref) {
+            Py_DECREF(v);
+        }
         return -1;
     }
-
     if (i >= 0)
         ((unsigned long *)ap->ob_item)[i] = x;
+
+    if (do_decref) {
+        Py_DECREF(v);
+    }
     return 0;
 }
 
@@ -448,25 +461,28 @@ static int
 QQ_setitem(arrayobject *ap, Py_ssize_t i, PyObject *v)
 {
     unsigned PY_LONG_LONG x;
-    if (PyLong_Check(v)) {
-        x = PyLong_AsUnsignedLongLong(v);
-        if (x == (unsigned PY_LONG_LONG) -1 && PyErr_Occurred())
+    int do_decref = 0; /* if nb_int was called */
+
+    if (!PyLong_Check(v)) {
+        v = get_int_unless_float(v);
+        if (NULL == v) {
             return -1;
+        }
+        do_decref = 1;
     }
-    else {
-        PY_LONG_LONG y;
-        if (!PyArg_Parse(v, "L;array item must be integer", &y))
-            return -1;
-        if (y < 0) {
-            PyErr_SetString(PyExc_OverflowError,
-                "unsigned long long is less than minimum");
-            return -1;
+    x = PyLong_AsUnsignedLongLong(v);
+    if (x == (unsigned PY_LONG_LONG)-1 && PyErr_Occurred()) {
+        if (do_decref) {
+            Py_DECREF(v);
         }
-        x = (unsigned PY_LONG_LONG)y;
+        return -1;
     }
-
     if (i >= 0)
         ((unsigned PY_LONG_LONG *)ap->ob_item)[i] = x;
+
+    if (do_decref) {
+        Py_DECREF(v);
+    }
     return 0;
 }
 #endif