From 4729529e2c7b0d1a629c3749eb9e242ba5e208a3 Mon Sep 17 00:00:00 2001
From: "bazel.build machine account" <ci.bazel@gmail.com>
Date: Wed, 10 Apr 2024 14:57:16 -0400
Subject: [PATCH] [7.2.0] Support `key` callback in Starlark min/max builtins
 (#21960)

This is required by the language spec, but was not implemented in Bazel.

See https://github.com/bazelbuild/starlark/blob/master/spec.md#max

Fixes https://github.com/bazelbuild/bazel/issues/15022

Also take the opportunity to adjust sorted's signature for `key` to
match.

RELNOTES: Starlark `min` and `max` buitins now allow a `key` callback,
similarly to `sorted`.
PiperOrigin-RevId: 623547043
Change-Id: I71d44aa715793f9f2260f9b20b876694154ff352

Commit
https://github.com/bazelbuild/bazel/commit/cf666726d0f8d1e8a3fed504810f80969dee4a4b

Co-authored-by: Googler <arostovtsev@google.com>
---
 .../net/starlark/java/eval/MethodLibrary.java | 173 +++++++++++++++---
 .../starlark/java/eval/testdata/min_max.star  |  90 ++++++++-
 .../starlark/java/eval/testdata/sorted.star   |  97 ++++++----
 3 files changed, 293 insertions(+), 67 deletions(-)

diff --git a/src/main/java/net/starlark/java/eval/MethodLibrary.java b/src/main/java/net/starlark/java/eval/MethodLibrary.java
index 36990a8586c53b..f5c50fb35999fa 100644
--- a/src/main/java/net/starlark/java/eval/MethodLibrary.java
+++ b/src/main/java/net/starlark/java/eval/MethodLibrary.java
@@ -14,12 +14,17 @@
 
 package net.starlark.java.eval;
 
+import static com.google.common.collect.Streams.stream;
+import static java.util.Comparator.comparing;
+
 import com.google.common.base.Ascii;
+import com.google.common.base.Throwables;
 import com.google.common.collect.Ordering;
 import java.util.Arrays;
 import java.util.Comparator;
 import java.util.Iterator;
 import java.util.NoSuchElementException;
+import java.util.Optional;
 import net.starlark.java.annot.Param;
 import net.starlark.java.annot.ParamType;
 import net.starlark.java.annot.StarlarkBuiltin;
@@ -31,39 +36,98 @@ class MethodLibrary {
   @StarlarkMethod(
       name = "min",
       doc =
-          "Returns the smallest one of all given arguments. "
-              + "If only one argument is provided, it must be a non-empty iterable. "
-              + "It is an error if elements are not comparable (for example int with string), "
-              + "or if no arguments are given. "
-              + "<pre class=\"language-python\">min(2, 5, 4) == 2\n"
-              + "min([5, 6, 3]) == 3</pre>",
-      extraPositionals = @Param(name = "args", doc = "The elements to be checked."))
-  public Object min(Sequence<?> args) throws EvalException {
-    return findExtreme(args, Starlark.ORDERING.reverse());
+          "Returns the smallest one of all given arguments. If only one positional argument is"
+              + " provided, it must be a non-empty iterable. It is an error if elements are not"
+              + " comparable (for example int with string), or if no arguments are given."
+              + "<pre class=\"language-python\">\n" //
+              + "min(2, 5, 4) == 2\n"
+              + "min([5, 6, 3]) == 3\n"
+              + "min(\"six\", \"three\", \"four\", key = len) == \"six\"  # the shortest\n"
+              + "min([2, -2, -1, 1], key = abs) == -1  # the first encountered with minimal key"
+              + " value\n"
+              + "</pre>",
+      extraPositionals = @Param(name = "args", doc = "The elements to be checked."),
+      parameters = {
+        @Param(
+            name = "key",
+            named = true,
+            positional = false,
+            allowedTypes = {
+              @ParamType(type = StarlarkCallable.class),
+              @ParamType(type = NoneType.class),
+            },
+            doc = "An optional function applied to each element before comparison.",
+            defaultValue = "None")
+      },
+      useStarlarkThread = true)
+  public Object min(Object key, Sequence<?> args, StarlarkThread thread)
+      throws EvalException, InterruptedException {
+    return findExtreme(
+        args,
+        Starlark.toJavaOptional(key, StarlarkCallable.class),
+        Starlark.ORDERING.reverse(),
+        thread);
   }
 
   @StarlarkMethod(
       name = "max",
       doc =
-          "Returns the largest one of all given arguments. "
-              + "If only one argument is provided, it must be a non-empty iterable."
-              + "It is an error if elements are not comparable (for example int with string), "
-              + "or if no arguments are given. "
-              + "<pre class=\"language-python\">max(2, 5, 4) == 5\n"
-              + "max([5, 6, 3]) == 6</pre>",
-      extraPositionals = @Param(name = "args", doc = "The elements to be checked."))
-  public Object max(Sequence<?> args) throws EvalException {
-    return findExtreme(args, Starlark.ORDERING);
+          "Returns the largest one of all given arguments. If only one positional argument is"
+              + " provided, it must be a non-empty iterable.It is an error if elements are not"
+              + " comparable (for example int with string), or if no arguments are given."
+              + "<pre class=\"language-python\">\n" //
+              + "max(2, 5, 4) == 5\n"
+              + "max([5, 6, 3]) == 6\n"
+              + "max(\"two\", \"three\", \"four\", key = len) ==\"three\"  # the longest\n"
+              + "max([1, -1, -2, 2], key = abs) == -2  # the first encountered with maximal key"
+              + " value\n"
+              + "</pre>",
+      extraPositionals = @Param(name = "args", doc = "The elements to be checked."),
+      parameters = {
+        @Param(
+            name = "key",
+            named = true,
+            positional = false,
+            allowedTypes = {
+              @ParamType(type = StarlarkCallable.class),
+              @ParamType(type = NoneType.class),
+            },
+            doc = "An optional function applied to each element before comparison.",
+            defaultValue = "None")
+      },
+      useStarlarkThread = true)
+  public Object max(Object key, Sequence<?> args, StarlarkThread thread)
+      throws EvalException, InterruptedException {
+    return findExtreme(
+        args, Starlark.toJavaOptional(key, StarlarkCallable.class), Starlark.ORDERING, thread);
   }
 
   /** Returns the maximum element from this list, as determined by maxOrdering. */
-  private static Object findExtreme(Sequence<?> args, Ordering<Object> maxOrdering)
-      throws EvalException {
+  private static Object findExtreme(
+      Sequence<?> args,
+      Optional<StarlarkCallable> keyFn,
+      Ordering<Object> maxOrdering,
+      StarlarkThread thread)
+      throws EvalException, InterruptedException {
     // Args can either be a list of items to compare, or a singleton list whose element is an
     // iterable of items to compare. In either case, there must be at least one item to compare.
     Iterable<?> items = (args.size() == 1) ? Starlark.toIterable(args.get(0)) : args;
     try {
-      return maxOrdering.max(items);
+      if (keyFn.isPresent()) {
+        try {
+          return stream(items)
+              .map(value -> ValueWithComparisonKey.make(value, keyFn.get(), thread))
+              .max(comparing(ValueWithComparisonKey::getComparisonKey, maxOrdering))
+              .get()
+              .getValue();
+        } catch (ValueWithComparisonKey.KeyCallException ex) {
+          Throwables.throwIfInstanceOf(ex.getCause(), EvalException.class);
+          Throwables.throwIfInstanceOf(ex.getCause(), InterruptedException.class);
+          throw new AssertionError("Got invalid ValueWithComparisonKey.KeyCallException", ex);
+        }
+      } else {
+        return maxOrdering.max(items);
+      }
     } catch (ClassCastException ex) {
       throw new EvalException(ex.getMessage()); // e.g. unsupported comparison: int <=> string
     } catch (NoSuchElementException ex) {
@@ -71,6 +135,52 @@ private static Object findExtreme(Sequence<?> args, Ordering<Object> maxOrdering
     }
   }
 
+  /**
+   * Original value decorated with its comparison key; storing the comparison key alongside the
+   * value ensures that we call the comparison key computation function only once per original value
+   * (which is important in case the function has side effects).
+   */
+  private static final class ValueWithComparisonKey {
+    private final Object value;
+    private final Object comparisonKey;
+
+    private ValueWithComparisonKey(Object value, Object comparisonKey) {
+      this.value = value;
+      this.comparisonKey = comparisonKey;
+    }
+
+    /**
+     * @throws KeyCallException wrapping the exception thrown by the underlying {@link
+     *     Starlark#fastcall} call if it threw.
+     */
+    static ValueWithComparisonKey make(
+        Object value, StarlarkCallable keyFn, StarlarkThread thread) {
+      Object[] positional = {value};
+      Object[] named = {};
+      try {
+        return new ValueWithComparisonKey(
+            value, Starlark.fastcall(thread, keyFn, positional, named));
+      } catch (EvalException | InterruptedException ex) {
+        throw new KeyCallException(ex);
+      }
+    }
+
+    Object getValue() {
+      return value;
+    }
+
+    Object getComparisonKey() {
+      return comparisonKey;
+    }
+
+    /** An unchecked exception wrapping an exception thrown by {@link Starlark#fastcall}. */
+    private static final class KeyCallException extends RuntimeException {
+      KeyCallException(Exception cause) {
+        super(cause);
+      }
+    }
+  }
+
   @StarlarkMethod(
       name = "abs",
       doc =
@@ -140,16 +250,24 @@ private static boolean hasElementWithBooleanValue(Object seq, boolean value)
               + " using x < y. The elements are sorted into ascending order, unless the reverse"
               + " argument is True, in which case the order is descending.\n"
               + " Sorting is stable: elements that compare equal retain their original relative"
-              + " order.\n"
-              + "<pre class=\"language-python\">sorted([3, 5, 4]) == [3, 4, 5]</pre>",
+              + " order.\n" //
+              + "<pre class=\"language-python\">\n" //
+              + "sorted([3, 5, 4]) == [3, 4, 5]\n" //
+              + "sorted([3, 5, 4], reverse = True) == [5, 4, 3]\n" //
+              + "sorted([\"two\", \"three\", \"four\"], key = len) == [\"two\", \"four\","
+              + " \"three\"]  # sort by length\n" //
+              + "</pre>",
       parameters = {
         @Param(name = "iterable", doc = "The iterable sequence to sort."),
         @Param(
             name = "key",
-            doc = "An optional function applied to each element before comparison.",
             named = true,
-            defaultValue = "None",
-            positional = false),
+            allowedTypes = {
+              @ParamType(type = StarlarkCallable.class),
+              @ParamType(type = NoneType.class),
+            },
+            doc = "An optional function applied to each element before comparison.",
+            defaultValue = "None"),
         @Param(
             name = "reverse",
             doc = "Return results in descending order.",
@@ -177,9 +295,6 @@ public StarlarkList<?> sorted(
     // The user provided a key function.
     // We must call it exactly once per element, in order,
     // so use the decorate/sort/undecorate pattern.
-    if (!(key instanceof StarlarkCallable)) {
-      throw Starlark.errorf("for key, got %s, want callable", Starlark.type(key));
-    }
     StarlarkCallable keyfn = (StarlarkCallable) key;
 
     // decorate
diff --git a/src/test/java/net/starlark/java/eval/testdata/min_max.star b/src/test/java/net/starlark/java/eval/testdata/min_max.star
index 2288bcdb11a202..d9b6a5deec888d 100644
--- a/src/test/java/net/starlark/java/eval/testdata/min_max.star
+++ b/src/test/java/net/starlark/java/eval/testdata/min_max.star
@@ -8,8 +8,8 @@ assert_eq(min([1, 2], [3]), [1, 2])
 assert_eq(min([1, 5], [1, 6], [2, 4], [0, 6]), [0, 6])
 assert_eq(min([-1]), -1)
 assert_eq(min([5, 2, 3]), 2)
-assert_eq(min({1: 2, -1: 3}), -1)
-assert_eq(min({2: None}), 2)
+assert_eq(min({1: 2, -1: 3}), -1)  # a single dict argument is treated as its sequence of keys
+assert_eq(min({2: None}), 2)  # a single dict argument is treated as its sequence of keys
 assert_eq(min(-1, 2), -1)
 assert_eq(min(5, 2, 3), 2)
 assert_eq(min(1, 1, 1, 1, 1, 1), 1)
@@ -21,6 +21,53 @@ assert_fails(lambda: min([]), "expected at least one item")
 assert_fails(lambda: min(1, "2", True), "unsupported comparison: int <=> string")
 assert_fails(lambda: min([1, "2", True]), "unsupported comparison: int <=> string")
 
+# min with key
+assert_eq(min("aBcDeFXyZ".elems(), key = lambda s: s.upper()), "a")
+assert_eq(min("test", "xyz", key = len), "xyz")
+assert_eq(min([4, 5], [1], key = lambda x: x), [1])
+assert_eq(min([1, 2], [3], key = lambda x: x), [1, 2])
+assert_eq(min([1, 5], [1, 6], [2, 4], [0, 6], key = lambda x: x), [0, 6])
+assert_eq(min([1, 5], [1, 6], [2, 4], [0, 6], key = lambda x: x[1]), [2, 4])
+assert_eq(min([-1], key = lambda x: x), -1)
+assert_eq(min([5, 2, 3], key = lambda x: x), 2)
+assert_eq(min({1: 2, -1: 3}, key = lambda x: x), -1)  # a single dict argument is treated as its sequence of keys
+assert_eq(min({2: None}, key = lambda x: x), 2)  # a single dict argument is treated as its sequence of keys
+assert_eq(min(-1, 2, key = lambda x: x), -1)
+assert_eq(min(5, 2, 3, key = lambda x: x), 2)
+assert_eq(min(1, 1, 1, 1, 1, 1, key = lambda x: -x), 1)
+assert_eq(min([1, 1, 1, 1, 1, 1], key = lambda x: -x), 1)
+assert_fails(lambda: min(1, key = lambda x: x), "type 'int' is not iterable")
+assert_fails(lambda: min(key = lambda x: x), "expected at least one item")
+assert_fails(lambda: min([], key = lambda x: x), "expected at least one item")
+assert_fails(lambda: min([1], ["2"], [True], key = lambda x: x[0]), "unsupported comparison: (int <=> string|string <=> int)")
+assert_fails(lambda: min([[1], ["2"], [True]], key = lambda x: x[0]), "unsupported comparison: (int <=> string|string <=> int)")
+
+# verify min with key chooses first value with minimal key
+assert_eq(min(1, -1, -2, 2, key = abs), 1)
+assert_eq(min([1, -1, -2, 2], key = abs), 1)
+
+# min with failing key
+assert_fails(lambda: min(0, 1, 2, 3, 4, key = lambda x: "foo".elems()[x]), "index out of range \\(index is 3, but sequence has 3 elements\\)")
+assert_fails(lambda: min([0, 1, 2, 3, 4], key = lambda x: "foo".elems()[x]), "index out of range \\(index is 3, but sequence has 3 elements\\)")
+
+# min with non-callable key
+assert_fails(lambda: min(1, 2, 3, key = "hello"), "parameter 'key' got value of type 'string', want 'callable or NoneType'")
+assert_fails(lambda: min([1, 2, 3], key = "hello"), "parameter 'key' got value of type 'string', want 'callable or NoneType'")
+
+# verify min with key invokes key callback exactly once per item
+def make_counting_identity():
+    call_count = {}
+
+    def counting_identity(x):
+        call_count[x] = call_count.get(x, 0) + 1
+        return x
+
+    return counting_identity, call_count
+
+min_counting_identity, min_call_count = make_counting_identity()
+assert_eq(min("min".elems(), key = min_counting_identity), "i")
+assert_eq(min_call_count, {"m": 1, "i": 1, "n": 1})
+
 # max
 assert_eq(max("abcdefxyz".elems()), "z")
 assert_eq(max("test", "xyz"), "xyz")
@@ -28,8 +75,8 @@ assert_eq(max("test", "xyz"), "xyz")
 assert_eq(max([1, 2], [5]), [5])
 assert_eq(max([-1]), -1)
 assert_eq(max([5, 2, 3]), 5)
-assert_eq(max({1: 2, -1: 3}), 1)
-assert_eq(max({2: None}), 2)
+assert_eq(max({1: 2, -1: 3}), 1)  # a single dict argument is treated as its sequence of keys
+assert_eq(max({2: None}), 2)  # a single dict argument is treated as its sequence of keys
 assert_eq(max(-1, 2), 2)
 assert_eq(max(5, 2, 3), 5)
 assert_eq(max(1, 1, 1, 1, 1, 1), 1)
@@ -40,3 +87,38 @@ assert_fails(lambda: max(), "expected at least one item")
 assert_fails(lambda: max([]), "expected at least one item")
 assert_fails(lambda: max(1, "2", True), "unsupported comparison: int <=> string")
 assert_fails(lambda: max([1, "2", True]), "unsupported comparison: int <=> string")
+
+# max with key
+assert_eq(max("aBcDeFXyZ".elems(), key = lambda s: s.lower()), "Z")
+assert_eq(max("test", "xyz", key = len), "test")
+assert_eq(max([1, 2], [5], key = lambda x: x), [5])
+assert_eq(max([-1], key = lambda x: x), -1)
+assert_eq(max([5, 2, 3], key = lambda x: x), 5)
+assert_eq(max({1: 2, -1: 3}, key = lambda x: x), 1)  # a single dict argument is treated as its sequence of keys
+assert_eq(max({2: None}, key = lambda x: x), 2)  # a single dict argument is treated as its sequence of keys
+assert_eq(max(-1, 2, key = lambda x: x), 2)
+assert_eq(max(5, 2, 3, key = lambda x: x), 5)
+assert_eq(max(1, 1, 1, 1, 1, 1, key = lambda x: -x), 1)
+assert_eq(max([1, 1, 1, 1, 1, 1], key = lambda x: -x), 1)
+assert_fails(lambda: max(1, key = lambda x: x), "type 'int' is not iterable")
+assert_fails(lambda: max(key = lambda x: x), "expected at least one item")
+assert_fails(lambda: max([], key = lambda x: x), "expected at least one item")
+assert_fails(lambda: max([1], ["2"], [True], key = lambda x: x[0]), "unsupported comparison: (int <=> string|string <=> int)")
+assert_fails(lambda: max([[1], ["2"], [True]], key = lambda x: x[0]), "unsupported comparison: (int <=> string|string <=> int)")
+
+# verify max with key chooses first value with minimal key
+assert_eq(max(1, -1, -2, 2, key = abs), -2)
+assert_eq(max([1, -1, -2, 2], key = abs), -2)
+
+# max with failing key
+assert_fails(lambda: max(0, 1, 2, 3, 4, key = lambda i: "xyz".elems()[i]), "index out of range \\(index is 3, but sequence has 3 elements\\)")
+assert_fails(lambda: max([0, 1, 2, 3, 4], key = lambda i: "xyz".elems()[i]), "index out of range \\(index is 3, but sequence has 3 elements\\)")
+
+# max with non-callable key
+assert_fails(lambda: max(1, 2, 3, key = "hello"), "parameter 'key' got value of type 'string', want 'callable or NoneType'")
+assert_fails(lambda: max([1, 2, 3], key = "hello"), "parameter 'key' got value of type 'string', want 'callable or NoneType'")
+
+# verify max with key invokes key callback exactly once per item
+max_counting_identity, max_call_count = make_counting_identity()
+assert_eq(max("max".elems(), key = max_counting_identity), "x")
+assert_eq(max_call_count, {"m": 1, "a": 1, "x": 1})
diff --git a/src/test/java/net/starlark/java/eval/testdata/sorted.star b/src/test/java/net/starlark/java/eval/testdata/sorted.star
index aa6d187c9b54d2..d7446794b5a5aa 100644
--- a/src/test/java/net/starlark/java/eval/testdata/sorted.star
+++ b/src/test/java/net/starlark/java/eval/testdata/sorted.star
@@ -11,54 +11,79 @@ assert_eq(sorted([{}]), [{}])
 
 assert_eq(sorted([42, 123, 3]), [3, 42, 123])
 assert_eq(sorted(["wiz", "foo", "bar"]), ["bar", "foo", "wiz"])
-assert_eq(sorted([42, 123, 3], reverse=True), [123, 42, 3])
-assert_eq(sorted(["wiz", "foo", "bar"], reverse=True), ["wiz", "foo", "bar"])
-assert_eq(sorted(list({"a": 1, "b": 2})), ['a', 'b'])
+assert_eq(sorted([42, 123, 3], reverse = True), [123, 42, 3])
+assert_eq(sorted(["wiz", "foo", "bar"], reverse = True), ["wiz", "foo", "bar"])
+assert_eq(sorted(list({"a": 1, "b": 2})), ["a", "b"])
 
 def f(x):
-  return x[0]
+    return x[0]
+
 pairs = [(4, 0), (3, 1), (4, 2), (2, 3), (3, 4), (1, 5), (2, 6), (3, 7)]
 
-assert_eq(sorted(pairs, key=f),
-         [(1, 5),
-          (2, 3), (2, 6),
-          (3, 1), (3, 4), (3, 7),
-          (4, 0), (4, 2)])
-
-assert_eq(sorted(["two", "three", "four"], key=len),
-         ["two", "four", "three"])
-assert_eq(sorted(["two", "three", "four"], key=len, reverse=True),
-         ["three", "four", "two"])
-assert_eq(sorted([[1, 5], [0, 10], [4]], key=max),
-         [[4], [1, 5], [0, 10]])
-assert_eq(sorted([[1, 5], [0, 10], [4]], key=min, reverse=True),
-         [[4], [1, 5], [0, 10]])
-assert_eq(sorted([[2, 6, 1], [5, 2, 1], [1, 4, 2]], key=sorted),
-         [[1, 4, 2], [5, 2, 1], [2, 6, 1]])
+assert_eq(
+    sorted(pairs, key = f),
+    [
+        (1, 5),
+        (2, 3),
+        (2, 6),
+        (3, 1),
+        (3, 4),
+        (3, 7),
+        (4, 0),
+        (4, 2),
+    ],
+)
+
+assert_eq(
+    sorted(["two", "three", "four"], key = len),
+    ["two", "four", "three"],
+)
+assert_eq(
+    sorted(["two", "three", "four"], key = len, reverse = True),
+    ["three", "four", "two"],
+)
+assert_eq(
+    sorted([[1, 5], [0, 10], [4]], key = max),
+    [[4], [1, 5], [0, 10]],
+)
+assert_eq(
+    sorted([[1, 5], [0, 10], [4]], key = min, reverse = True),
+    [[4], [1, 5], [0, 10]],
+)
+assert_eq(
+    sorted([[2, 6, 1], [5, 2, 1], [1, 4, 2]], key = sorted),
+    [[1, 4, 2], [5, 2, 1], [2, 6, 1]],
+)
 
 # case-insensitive sort
-def lower(str): return str.lower()
+def lower(str):
+    return str.lower()
+
 assert_eq(sorted(["a", "B", "c"]), ["B", "a", "c"])
-assert_eq(sorted(["a", "B", "c"], reverse=True), ["c", "a", "B"])
-assert_eq(sorted(["a", "B", "c"], key=lower), ["a", "B", "c"])
-assert_eq(sorted(["a", "B", "c"], key=lower, reverse=True), ["c", "B", "a"])
+assert_eq(sorted(["a", "B", "c"], reverse = True), ["c", "a", "B"])
+assert_eq(sorted(["a", "B", "c"], key = lower), ["a", "B", "c"])
+assert_eq(sorted(["a", "B", "c"], key = lower, reverse = True), ["c", "B", "a"])
 
 # NaN values compare greater even than +Inf.
 # (By contrast, Python's 'sorted' sorts the numbers between the NaN values,
 # producing a pairwise but not transitively nondescending result.)
 # See for example https://bugs.python.org/issue36095.
-n=float('nan')
-assert_eq(str(sorted([7, 4, n, 19, 2, 6, 3, 87, 17, n, 5, 31, 12, 6, 4, 4, 2, n, 2, 1])),
-          "[1, 2, 2, 2, 3, 4, 4, 4, 5, 6, 6, 7, 12, 17, 19, 31, 87, nan, nan, nan]")
+n = float("nan")
+assert_eq(
+    str(sorted([7, 4, n, 19, 2, 6, 3, 87, 17, n, 5, 31, 12, 6, 4, 4, 2, n, 2, 1])),
+    "[1, 2, 2, 2, 3, 4, 4, 4, 5, 6, 6, 7, 12, 17, 19, 31, 87, nan, nan, nan]",
+)
 assert_eq(str(sorted([7, 3, n, 1, 9])), "[1, 3, 7, 9, nan]")
-assert_eq(str(sorted([7, 3, n, 1, 9], reverse=True)), "[nan, 9, 7, 3, 1]")
+assert_eq(str(sorted([7, 3, n, 1, 9], reverse = True)), "[nan, 9, 7, 3, 1]")
 
 # The key function is called once per array element, in order
 acc = []
+
 def keyfn(k):
     acc.append(k)
     return k
-assert_eq(sorted([3, 1, 4, 1, 5, 9], key=keyfn), [1, 1, 3, 4, 5, 9])
+
+assert_eq(sorted([3, 1, 4, 1, 5, 9], key = keyfn), [1, 1, 3, 4, 5, 9])
 assert_eq(acc, [3, 1, 4, 1, 5, 9])
 
 # Ordered list comparison x < y finds the first unequal elements
@@ -66,8 +91,8 @@ assert_eq(acc, [3, 1, 4, 1, 5, 9])
 # Thus it is not an error if the lists contain equal elements
 # that do not support ordered comparison, such as dict or None.
 d1, d2 = {}, {}
-assert_([d1] <= [d1]) # same object
-assert_([d1] <= [d2]) # distinct objects
+assert_([d1] <= [d1])  # same object
+assert_([d1] <= [d2])  # distinct objects
 assert_([d1] <= [d1, d2])
 assert_([None] <= [None])
 
@@ -79,12 +104,16 @@ assert_([None] <= [None])
 # The first compares greater than the second.
 assert_eq(sorted(["�", "🌿"]), ["🌿", "�"])
 
-
 assert_(False < True)
 assert_fails(lambda: False < 1, "unsupported comparison: bool <=> int")
-assert_fails(lambda: [{1: None}] <= [{2:None}], "unsupported comparison: dict <=> dict")
+assert_fails(lambda: [{1: None}] <= [{2: None}], "unsupported comparison: dict <=> dict")
 assert_fails(lambda: None <= None, "unsupported comparison: NoneType <=> NoneType")
 assert_fails(lambda: sorted(1), "got value of type 'int', want 'iterable'")
 assert_fails(lambda: sorted([1, 2, None, 3]), "unsupported comparison: NoneType <=> int")
 assert_fails(lambda: sorted([1, "one"]), "unsupported comparison: string <=> int")
-assert_fails(lambda: sorted([1, 2, 3], key=1), "for key, got int, want callable")
+
+# non-callable key
+assert_fails(lambda: sorted([1, 2, 3], key = 1), "parameter 'key' got value of type 'int', want 'callable or NoneType'")
+
+# failing key computation
+assert_fails(lambda: sorted([1, 2, 3], key = lambda i: "abc".elems()[i]), "index out of range \\(index is 3, but sequence has 3 elements\\)")