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

1 files changed, 48 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 *