Merge pull request RustPython#3695 from rebunto/optional-key-argument-for-_bisect-functions

Optional key argument for bisect functions

Author: youknowone

Lib/bisect.py

@@ -1,29 +1,26 @@
 """Bisection algorithms."""
 
-def insort_right(a, x, lo=0, hi=None):
+
+def insort_right(a, x, lo=0, hi=None, *, key=None):
     """Insert item x in list a, and keep it sorted assuming a is sorted.
 
     If x is already in a, insert it to the right of the rightmost x.
 
     Optional args lo (default 0) and hi (default len(a)) bound the
     slice of a to be searched.
     """
-
-    if lo < 0:
-        raise ValueError('lo must be non-negative')
-    if hi is None:
-        hi = len(a)
-    while lo < hi:
-        mid = (lo+hi)//2
-        if x < a[mid]: hi = mid
-        else: lo = mid+1
+    if key is None:
+        lo = bisect_right(a, x, lo, hi)
+    else:
+        lo = bisect_right(a, key(x), lo, hi, key=key)
     a.insert(lo, x)
 
-def bisect_right(a, x, lo=0, hi=None):
+
+def bisect_right(a, x, lo=0, hi=None, *, key=None):
     """Return the index where to insert item x in list a, assuming a is sorted.
 
     The return value i is such that all e in a[:i] have e <= x, and all e in
-    a[i:] have e > x.  So if x already appears in the list, a.insert(x) will
+    a[i:] have e > x.  So if x already appears in the list, a.insert(i, x) will
     insert just after the rightmost x already there.
 
     Optional args lo (default 0) and hi (default len(a)) bound the
@@ -34,13 +31,26 @@ def bisect_right(a, x, lo=0, hi=None):
         raise ValueError('lo must be non-negative')
     if hi is None:
         hi = len(a)
-    while lo < hi:
-        mid = (lo+hi)//2
-        if x < a[mid]: hi = mid
-        else: lo = mid+1
+    # Note, the comparison uses "<" to match the
+    # __lt__() logic in list.sort() and in heapq.
+    if key is None:
+        while lo < hi:
+            mid = (lo + hi) // 2
+            if x < a[mid]:
+                hi = mid
+            else:
+                lo = mid + 1
+    else:
+        while lo < hi:
+            mid = (lo + hi) // 2
+            if x < key(a[mid]):
+                hi = mid
+            else:
+                lo = mid + 1
     return lo
 
-def insort_left(a, x, lo=0, hi=None):
+
+def insort_left(a, x, lo=0, hi=None, *, key=None):
     """Insert item x in list a, and keep it sorted assuming a is sorted.
 
     If x is already in a, insert it to the left of the leftmost x.
@@ -49,22 +59,17 @@ def insort_left(a, x, lo=0, hi=None):
     slice of a to be searched.
     """
 
-    if lo < 0:
-        raise ValueError('lo must be non-negative')
-    if hi is None:
-        hi = len(a)
-    while lo < hi:
-        mid = (lo+hi)//2
-        if a[mid] < x: lo = mid+1
-        else: hi = mid
+    if key is None:
+        lo = bisect_left(a, x, lo, hi)
+    else:
+        lo = bisect_left(a, key(x), lo, hi, key=key)
     a.insert(lo, x)
 
-
-def bisect_left(a, x, lo=0, hi=None):
+def bisect_left(a, x, lo=0, hi=None, *, key=None):
     """Return the index where to insert item x in list a, assuming a is sorted.
 
     The return value i is such that all e in a[:i] have e < x, and all e in
-    a[i:] have e >= x.  So if x already appears in the list, a.insert(x) will
+    a[i:] have e >= x.  So if x already appears in the list, a.insert(i, x) will
     insert just before the leftmost x already there.
 
     Optional args lo (default 0) and hi (default len(a)) bound the
@@ -75,17 +80,31 @@ def bisect_left(a, x, lo=0, hi=None):
         raise ValueError('lo must be non-negative')
     if hi is None:
         hi = len(a)
-    while lo < hi:
-        mid = (lo+hi)//2
-        if a[mid] < x: lo = mid+1
-        else: hi = mid
+    # Note, the comparison uses "<" to match the
+    # __lt__() logic in list.sort() and in heapq.
+    if key is None:
+        while lo < hi:
+            mid = (lo + hi) // 2
+            if a[mid] < x:
+                lo = mid + 1
+            else:
+                hi = mid
+    else:
+        while lo < hi:
+            mid = (lo + hi) // 2
+            if key(a[mid]) < x:
+                lo = mid + 1
+            else:
+                hi = mid
     return lo
 
-# Overwrite above definitions with a fast Rust implementation
+
+# Overwrite above definitions with a fast C implementation
 try:
     from _bisect import *
 except ImportError:
     pass
 
-bisect = bisect_right   # backward compatibility
-insort = insort_right   # backward compatibility
+# Create aliases
+bisect = bisect_right
+insort = insort_right

Lib/test/test_bisect.py

@@ -1,9 +1,9 @@
 import sys
 import unittest
-from test import support
 from test.support import import_helper
 from collections import UserList
 
+
 py_bisect = import_helper.import_fresh_module('bisect', blocked=['_bisect'])
 c_bisect = import_helper.import_fresh_module('bisect', fresh=['bisect'])
 
@@ -200,6 +200,63 @@ def test_keyword_args(self):
         self.module.insort(a=data, x=25, lo=1, hi=3)
         self.assertEqual(data, [10, 20, 25, 25, 25, 30, 40, 50])
 
+    def test_lookups_with_key_function(self):
+        mod = self.module
+
+        # Invariant: Index with a keyfunc on an array
+        # should match the index on an array where
+        # key function has already been applied.
+
+        keyfunc = abs
+        arr = sorted([2, -4, 6, 8, -10], key=keyfunc)
+        precomputed_arr = list(map(keyfunc, arr))
+        for x in precomputed_arr:
+            self.assertEqual(
+                mod.bisect_left(arr, x, key=keyfunc),
+                mod.bisect_left(precomputed_arr, x)
+            )
+            self.assertEqual(
+                mod.bisect_right(arr, x, key=keyfunc),
+                mod.bisect_right(precomputed_arr, x)
+            )
+
+        keyfunc = str.casefold
+        arr = sorted('aBcDeEfgHhiIiij', key=keyfunc)
+        precomputed_arr = list(map(keyfunc, arr))
+        for x in precomputed_arr:
+            self.assertEqual(
+                mod.bisect_left(arr, x, key=keyfunc),
+                mod.bisect_left(precomputed_arr, x)
+            )
+            self.assertEqual(
+                mod.bisect_right(arr, x, key=keyfunc),
+                mod.bisect_right(precomputed_arr, x)
+            )
+
+    def test_insort(self):
+        from random import shuffle
+        mod = self.module
+
+        # Invariant:  As random elements are inserted in
+        # a target list, the targetlist remains sorted.
+        keyfunc = abs
+        data = list(range(-10, 11)) + list(range(-20, 20, 2))
+        shuffle(data)
+        target = []
+        for x in data:
+            mod.insort_left(target, x, key=keyfunc)
+            self.assertEqual(
+                sorted(target, key=keyfunc),
+                target
+            )
+        target = []
+        for x in data:
+            mod.insort_right(target, x, key=keyfunc)
+            self.assertEqual(
+                sorted(target, key=keyfunc),
+                target
+            )
+
 class TestBisectPython(TestBisect, unittest.TestCase):
     module = py_bisect
 

stdlib/src/bisect.rs

@@ -14,6 +14,8 @@ mod _bisect {
         lo: OptionalArg<PyObjectRef>,
         #[pyarg(any, optional)]
         hi: OptionalArg<PyObjectRef>,
+        #[pyarg(named, default)]
+        key: Option<PyObjectRef>,
     }
 
     // Handles objects that implement __index__ and makes sure index fits in needed isize.
@@ -66,17 +68,21 @@ mod _bisect {
     #[inline]
     #[pyfunction]
     fn bisect_left(
-        BisectArgs { a, x, lo, hi }: BisectArgs,
+        BisectArgs { a, x, lo, hi, key }: BisectArgs,
         vm: &VirtualMachine,
     ) -> PyResult<usize> {
         let (mut lo, mut hi) = as_usize(lo, hi, a.length(vm)?, vm)?;
 
         while lo < hi {
             // Handles issue 13496.
             let mid = (lo + hi) / 2;
-            if a.get_item(&mid, vm)?
-                .rich_compare_bool(&x, PyComparisonOp::Lt, vm)?
-            {
+            let a_mid = a.get_item(&mid, vm)?;
+            let comp = if let Some(ref key) = key {
+                vm.invoke(key, (a_mid,))?
+            } else {
+                a_mid
+            };
+            if comp.rich_compare_bool(&x, PyComparisonOp::Lt, vm)? {
                 lo = mid + 1;
             } else {
                 hi = mid;
@@ -96,15 +102,21 @@ mod _bisect {
     #[inline]
     #[pyfunction]
     fn bisect_right(
-        BisectArgs { a, x, lo, hi }: BisectArgs,
+        BisectArgs { a, x, lo, hi, key }: BisectArgs,
         vm: &VirtualMachine,
     ) -> PyResult<usize> {
         let (mut lo, mut hi) = as_usize(lo, hi, a.length(vm)?, vm)?;
 
         while lo < hi {
             // Handles issue 13496.
             let mid = (lo + hi) / 2;
-            if x.rich_compare_bool(&*a.get_item(&mid, vm)?, PyComparisonOp::Lt, vm)? {
+            let a_mid = a.get_item(&mid, vm)?;
+            let comp = if let Some(ref key) = key {
+                vm.invoke(key, (a_mid,))?
+            } else {
+                a_mid
+            };
+            if x.rich_compare_bool(&*comp, PyComparisonOp::Lt, vm)? {
                 hi = mid;
             } else {
                 lo = mid + 1;
@@ -120,13 +132,19 @@ mod _bisect {
     /// Optional args lo (default 0) and hi (default len(a)) bound the
     /// slice of a to be searched.
     #[pyfunction]
-    fn insort_left(BisectArgs { a, x, lo, hi }: BisectArgs, vm: &VirtualMachine) -> PyResult {
+    fn insort_left(BisectArgs { a, x, lo, hi, key }: BisectArgs, vm: &VirtualMachine) -> PyResult {
+        let x = if let Some(ref key) = key {
+            vm.invoke(key, (x,))?
+        } else {
+            x
+        };
         let index = bisect_left(
             BisectArgs {
                 a: a.clone(),
                 x: x.clone(),
                 lo,
                 hi,
+                key,
             },
             vm,
         )?;
@@ -140,13 +158,19 @@ mod _bisect {
     /// Optional args lo (default 0) and hi (default len(a)) bound the
     /// slice of a to be searched
     #[pyfunction]
-    fn insort_right(BisectArgs { a, x, lo, hi }: BisectArgs, vm: &VirtualMachine) -> PyResult {
+    fn insort_right(BisectArgs { a, x, lo, hi, key }: BisectArgs, vm: &VirtualMachine) -> PyResult {
+        let x = if let Some(ref key) = key {
+            vm.invoke(key, (x,))?
+        } else {
+            x
+        };
         let index = bisect_right(
             BisectArgs {
                 a: a.clone(),
                 x: x.clone(),
                 lo,
                 hi,
+                key,
             },
             vm,
         )?;