Add key and reverse parameters to builtin sorted function
Implement:https://github.com/bazelbuild/starlark/issues/20#issuecomment-456647994
Closes #8881.
PiperOrigin-RevId: 258401453
diff --git a/src/main/java/com/google/devtools/build/lib/syntax/MethodLibrary.java b/src/main/java/com/google/devtools/build/lib/syntax/MethodLibrary.java
index c32b810..557fc60 100644
--- a/src/main/java/com/google/devtools/build/lib/syntax/MethodLibrary.java
+++ b/src/main/java/com/google/devtools/build/lib/syntax/MethodLibrary.java
@@ -36,6 +36,8 @@
import com.google.devtools.build.lib.syntax.SkylarkList.Tuple;
import java.util.ArrayDeque;
import java.util.ArrayList;
+import java.util.Collections;
+import java.util.Comparator;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
@@ -102,6 +104,22 @@
}
}
+ private static Object evalKeyFunc(
+ Object obj, Object key, Location loc, Environment env, FuncallExpression ast)
+ throws EvalException, InterruptedException {
+ checkValidKeyFunc(key, loc);
+ return ((StarlarkFunction) key)
+ .call(Collections.singletonList(obj), Collections.emptyMap(), ast, env);
+ }
+
+ private static void checkValidKeyFunc(Object key, Location loc) throws EvalException {
+ if (key instanceof StarlarkFunction) {
+ return;
+ }
+ throw new EvalException(
+ loc, Printer.format("%r object is not callable", EvalUtils.getDataTypeName(key)));
+ }
+
@SkylarkCallable(
name = "all",
doc =
@@ -170,17 +188,77 @@
type = Object.class,
doc = "This collection.",
// TODO(cparsons): This parameter should be positional-only.
- legacyNamed = true)
+ legacyNamed = true),
+ @Param(
+ name = "key",
+ doc = "An optional function applied to each element before comparison.",
+ named = true,
+ defaultValue = "None",
+ positional = false,
+ noneable = true),
+ @Param(
+ name = "reverse",
+ type = Boolean.class,
+ doc = "Return results in descending order.",
+ named = true,
+ defaultValue = "False",
+ positional = false,
+ noneable = true)
},
useLocation = true,
useEnvironment = true)
- public MutableList<?> sorted(Object self, Location loc, Environment env) throws EvalException {
- try {
- return MutableList.copyOf(
- env, EvalUtils.SKYLARK_COMPARATOR.sortedCopy(EvalUtils.toCollection(self, loc, env)));
- } catch (EvalUtils.ComparisonException e) {
- throw new EvalException(loc, e);
+ public MutableList<?> sorted(
+ Object self, final Object key, Boolean reverse, final Location loc, final Environment env)
+ throws EvalException, InterruptedException {
+
+ ArrayList list = new ArrayList(EvalUtils.toCollection(self, loc, env));
+ if (key == Runtime.NONE) {
+ try {
+ Collections.sort(list, EvalUtils.SKYLARK_COMPARATOR);
+ } catch (EvalUtils.ComparisonException e) {
+ throw new EvalException(loc, e);
+ }
+ } else {
+ checkValidKeyFunc(key, loc);
+
+ final FuncallExpression ast = new FuncallExpression(Identifier.of(""), ImmutableList.of());
+
+ class KeyComparator implements Comparator<Object> {
+ Exception e;
+
+ @Override
+ public int compare(Object x, Object y) {
+ try {
+ return EvalUtils.SKYLARK_COMPARATOR.compare(
+ evalKeyFunc(x, key, loc, env, ast), evalKeyFunc(y, key, loc, env, ast));
+ } catch (InterruptedException | EvalException e) {
+ if (this.e == null) {
+ this.e = e;
+ }
+ return 0;
+ }
+ }
+ }
+
+ KeyComparator comp = new KeyComparator();
+ try {
+ Collections.sort(list, comp);
+ } catch (EvalUtils.ComparisonException e) {
+ throw new EvalException(loc, e);
+ }
+
+ if (comp.e != null) {
+ if (comp.e instanceof InterruptedException) {
+ throw (InterruptedException) comp.e;
+ }
+ throw (EvalException) comp.e;
+ }
}
+
+ if (reverse) {
+ Collections.reverse(list);
+ }
+ return MutableList.wrapUnsafe(env, list);
}
@SkylarkCallable(
diff --git a/src/test/java/com/google/devtools/build/lib/syntax/MethodLibraryTest.java b/src/test/java/com/google/devtools/build/lib/syntax/MethodLibraryTest.java
index aabd297..c5c3706 100644
--- a/src/test/java/com/google/devtools/build/lib/syntax/MethodLibraryTest.java
+++ b/src/test/java/com/google/devtools/build/lib/syntax/MethodLibraryTest.java
@@ -249,6 +249,9 @@
.testEval("sorted([True, False, True])", "[False, True, True]")
.testEval("sorted(['a','x','b','z'])", "[\"a\", \"b\", \"x\", \"z\"]")
.testEval("sorted({1: True, 5: True, 4: False})", "[1, 4, 5]")
+ .testEval("sorted([3, 2, 1, 0], reverse=True)", "[3, 2, 1, 0]")
+ .testEval("sorted([[1], [], [1, 2]], key=len, reverse=True)", "[[1, 2], [1], []]")
+ .testEval("sorted([[0, 5], [4, 1], [1, 7]], key=max)", "[[4, 1], [0, 5], [1, 7]]")
.testIfExactError("Cannot compare function with function", "sorted([sorted, sorted])");
}
@@ -712,7 +715,7 @@
// Parameters which may be specified by keyword but are not explicitly 'named'.
.testStatement("all(elements=[True, True])", Boolean.TRUE)
.testStatement("any(elements=[True, False])", Boolean.TRUE)
- .testEval("sorted(self=[3, 0, 2])", "[0, 2, 3]")
+ .testEval("sorted(self=[3, 0, 2], key=None, reverse=False)", "[0, 2, 3]")
.testEval("reversed(sequence=[3, 2, 0])", "[0, 2, 3]")
.testEval("tuple(x=[1, 2])", "(1, 2)")
.testEval("list(x=(1, 2))", "[1, 2]")
diff --git a/src/test/starlark/testdata/sorted.sky b/src/test/starlark/testdata/sorted.sky
new file mode 100644
index 0000000..d7fa3ab
--- /dev/null
+++ b/src/test/starlark/testdata/sorted.sky
@@ -0,0 +1,37 @@
+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'])
+---
+
+def f(x):
+ 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]])
+
+---
+sorted(1) ### type 'int' is not a collection
+---
+sorted([1, 2, None, 3]) ### Cannot compare NoneType with int
+---
+sorted([1, "one"]) ### Cannot compare string with int
+---
+sorted([1, 2, 3], key=1) ### "int" object is not callable