bpo-4356: Add key function support to the bisect module (GH-20556)

2 files changed, 105 insertions, 18 deletions

diff --git a/Lib/bisect.py b/Lib/bisect.py
index 8336a4e..d37da74 100644
--- a/Lib/bisect.py
+++ b/Lib/bisect.py
@@ -1,6 +1,7 @@
 """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.
@@ -8,11 +9,14 @@ def insort_right(a, x, lo=0, hi=None):
     Optional args lo (default 0) and hi (default len(a)) bound the
     slice of a to be searched.
     """
-
-    lo = bisect_right(a, x, lo, hi)
+    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
@@ -27,14 +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
-        # Use __lt__ to match the logic in list.sort() and in heapq
-        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.
@@ -43,11 +59,13 @@ def insort_left(a, x, lo=0, hi=None):
     slice of a to be searched.
     """
 
-    lo = bisect_left(a, x, lo, hi)
+    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
@@ -62,13 +80,25 @@ 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
-        # Use __lt__ to match the logic in list.sort() and in heapq
-        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 C implementation
 try:
     from _bisect import *
diff --git a/Lib/test/test_bisect.py b/Lib/test/test_bisect.py
index fc7990d..20f8b9d 100644
--- a/Lib/test/test_bisect.py
+++ b/Lib/test/test_bisect.py
@@ -200,6 +200,63 @@ class TestBisect:
         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