From cc2c4eea609f45c9370091abdf37c88530e64b23 Mon Sep 17 00:00:00 2001 From: Marwan Tammam Date: Tue, 16 Jul 2019 11:03:48 -0700 Subject: [PATCH] Add key and reverse parameters to builtin sorted function Implement:https://github.com/bazelbuild/starlark/issues/20#issuecomment-456647994 Closes #8881. PiperOrigin-RevId: 258401453 --- .../build/lib/syntax/MethodLibrary.java | 92 +++++++++++++++++-- .../build/lib/syntax/MethodLibraryTest.java | 5 +- src/test/starlark/testdata/sorted.sky | 37 ++++++++ 3 files changed, 126 insertions(+), 8 deletions(-) create mode 100644 src/test/starlark/testdata/sorted.sky 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 c32b810c292f3e..557fc60e207293 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 findExtreme( } } + 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 @@ private static boolean hasElementWithBooleanValue( 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 { + 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 aabd29710fe339..c5c3706be9abed 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 @@ public void testListSort() throws Exception { .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 @@ public void testLegacyNamed() throws Exception { // 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 00000000000000..d7fa3ab2d6c4a9 --- /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