Issue #1766304: Optimize membership testing for ranges: 'n in range(...)'

mdickinson · mdickinson · commit 3e124ae739db · 2009-09-22T21:47:24.000Z
does an O(1) check, if n is an integer.  Non-integers aren't affected.
Thanks Robert Lehmann.
diff --git a/Doc/library/functions.rst b/Doc/library/functions.rst
@@ -925,6 +925,10 @@ are always available.  They are listed here in alphabetical order.
       >>> list(range(1, 0))
       []
 
+   .. versionchanged:: 3.2
+      Testing integers for membership takes constant time instead of
+      iterating through all items.
+
 
 .. function:: repr(object)
 
diff --git a/Lib/test/test_range.py b/Lib/test/test_range.py
@@ -78,6 +78,56 @@ def test_odd_bug(self):
         with self.assertRaises(TypeError):
             range([], 1, -1)
 
+    def test_types(self):
+        # Non-integer objects *equal* to any of the range's items are supposed
+        # to be contained in the range.
+        self.assertTrue(1.0 in range(3))
+        self.assertTrue(True in range(3))
+        self.assertTrue(1+0j in range(3))
+
+        class C1:
+            def __eq__(self, other): return True
+        self.assertTrue(C1() in range(3))
+
+        # Objects are never coerced into other types for comparison.
+        class C2:
+            def __int__(self): return 1
+            def __index__(self): return 1
+        self.assertFalse(C2() in range(3))
+        # ..except if explicitly told so.
+        self.assertTrue(int(C2()) in range(3))
+
+
+    def test_strided_limits(self):
+        r = range(0, 101, 2)
+        self.assertTrue(0 in r)
+        self.assertFalse(1 in r)
+        self.assertTrue(2 in r)
+        self.assertFalse(99 in r)
+        self.assertTrue(100 in r)
+        self.assertFalse(101 in r)
+
+        r = range(0, -20, -1)
+        self.assertTrue(0 in r)
+        self.assertTrue(-1 in r)
+        self.assertTrue(-19 in r)
+        self.assertFalse(-20 in r)
+
+        r = range(0, -20, -2)
+        self.assertTrue(-18 in r)
+        self.assertFalse(-19 in r)
+        self.assertFalse(-20 in r)
+
+    def test_empty(self):
+        r = range(0)
+        self.assertFalse(0 in r)
+        self.assertFalse(1 in r)
+
+        r = range(0, -10)
+        self.assertFalse(0 in r)
+        self.assertFalse(-1 in r)
+        self.assertFalse(1 in r)
+
 def test_main():
     test.support.run_unittest(RangeTest)
 
diff --git a/Misc/NEWS b/Misc/NEWS
@@ -12,6 +12,8 @@ What's New in Python 3.2 Alpha 1?
 Core and Builtins
 -----------------
 
+- Issue #1766304: Improve performance of membership tests on range objects.
+
 - Issue #6713: Improve performance of integer -> string conversions.
 
 - Issue #6846: Fix bug where bytearray.pop() returns negative integers.
diff --git a/Objects/rangeobject.c b/Objects/rangeobject.c
@@ -273,12 +273,69 @@ range_reduce(rangeobject *r, PyObject *args)
                          r->start, r->stop, r->step);
 }
 
+static int
+range_contains(rangeobject *r, PyObject *ob) {
+    if (PyLong_Check(ob)) {
+        int cmp1, cmp2, cmp3;
+        PyObject *tmp1 = NULL;
+        PyObject *tmp2 = NULL;
+        PyObject *zero = NULL;
+        int result = -1;
+
+        zero = PyLong_FromLong(0);
+        if (zero == NULL) /* MemoryError in int(0) */
+            goto end;
+
+        /* Check if the value can possibly be in the range. */
+
+        cmp1 = PyObject_RichCompareBool(r->step, zero, Py_GT);
+        if (cmp1 == -1)
+            goto end;
+        if (cmp1 == 1) { /* positive steps: start <= ob < stop */
+            cmp2 = PyObject_RichCompareBool(r->start, ob, Py_LE);
+            cmp3 = PyObject_RichCompareBool(ob, r->stop, Py_LT);
+        }
+        else { /* negative steps: stop < ob <= start */
+            cmp2 = PyObject_RichCompareBool(ob, r->start, Py_LE);
+            cmp3 = PyObject_RichCompareBool(r->stop, ob, Py_LT);
+        }
+
+        if (cmp2 == -1 || cmp3 == -1) /* TypeError */
+            goto end;
+        if (cmp2 == 0 || cmp3 == 0) { /* ob outside of range */
+            result = 0;
+            goto end;
+        }
+
+        /* Check that the stride does not invalidate ob's membership. */
+        tmp1 = PyNumber_Subtract(ob, r->start);
+        if (tmp1 == NULL)
+            goto end;
+        tmp2 = PyNumber_Remainder(tmp1, r->step);
+        if (tmp2 == NULL)
+            goto end;
+        /* result = (int(ob) - start % step) == 0 */
+        result = PyObject_RichCompareBool(tmp2, zero, Py_EQ);
+      end:
+        Py_XDECREF(tmp1);
+        Py_XDECREF(tmp2);
+        Py_XDECREF(zero);
+        return result;
+    }
+    /* Fall back to iterative search. */
+    return (int)_PySequence_IterSearch((PyObject*)r, ob,
+                                       PY_ITERSEARCH_CONTAINS);
+}
+
 static PySequenceMethods range_as_sequence = {
     (lenfunc)range_length,	/* sq_length */
     0,			/* sq_concat */
     0,			/* sq_repeat */
     (ssizeargfunc)range_item, /* sq_item */
     0,			/* sq_slice */
+    0, /* sq_ass_item */
+    0, /* sq_ass_slice */
+    (objobjproc)range_contains, /* sq_contains */
 };
 
 static PyObject * range_iter(PyObject *seq);
