[3.6] bpo-28298: make array 'Q', 'L' and 'I' accept big intables as elements (#579)

Author: orenmn

Lib/test/test_array.py

@@ -14,14 +14,6 @@
 import array
 from array import _array_reconstructor as array_reconstructor
 
-try:
- # Try to determine availability of long long independently
- # of the array module under test
- struct.calcsize('@q')
- have_long_long = True
-except struct.error:
- have_long_long = False
-
 sizeof_wchar = array.array('u').itemsize
 
 
@@ -32,9 +24,7 @@ class ArraySubclassWithKwargs(array.array):
  def __init__(self, typecode, newarg=None):
  array.array.__init__(self)
 
-typecodes = "ubBhHiIlLfd"
-if have_long_long:
- typecodes += 'qQ'
+typecodes = 'ubBhHiIlLfdqQ'
 
 class MiscTest(unittest.TestCase):
 
@@ -1240,7 +1230,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]
@@ -1251,8 +1260,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]
@@ -1263,6 +1273,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)
@@ -1314,12 +1325,10 @@ class UnsignedLongTest(UnsignedNumberTest, unittest.TestCase):
  typecode = 'L'
  minitemsize = 4
 
-@unittest.skipIf(not have_long_long, 'need long long support')
 class LongLongTest(SignedNumberTest, unittest.TestCase):
  typecode = 'q'
  minitemsize = 8
 
-@unittest.skipIf(not have_long_long, 'need long long support')
 class UnsignedLongLongTest(UnsignedNumberTest, unittest.TestCase):
  typecode = 'Q'
  minitemsize = 8

Misc/NEWS

@@ -16,6 +16,9 @@ Core and Builtins
 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-28231: The zipfile module now accepts path-like objects for external
  paths.
 

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;
 }
 
@@ -446,25 +459,28 @@ static int
 QQ_setitem(arrayobject *ap, Py_ssize_t i, PyObject *v)
 {
  unsigned long long x;
- if (PyLong_Check(v)) {
- x = PyLong_AsUnsignedLongLong(v);
- if (x == (unsigned 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 {
- 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 long long)-1 && PyErr_Occurred()) {
+ if (do_decref) {
+ Py_DECREF(v);
  }
- x = (unsigned long long)y;
+ return -1;
  }
-
  if (i >= 0)
  ((unsigned long long *)ap->ob_item)[i] = x;
+
+ if (do_decref) {
+ Py_DECREF(v);
+ }
  return 0;
 }