Skylark: Allow list slices

--
MOS_MIGRATED_REVID=88727892

4 files changed, 49 insertions, 22 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/packages/MethodLibrary.java b/src/main/java/com/google/devtools/build/lib/packages/MethodLibrary.java
index 8035d08ffa..bee5ca17b6 100644
--- a/src/main/java/com/google/devtools/build/lib/packages/MethodLibrary.java
+++ b/src/main/java/com/google/devtools/build/lib/packages/MethodLibrary.java
@@ -67,18 +67,18 @@ public class MethodLibrary {
   // Convert string index in the same way Python does.
   // If index is negative, starts from the end.
   // If index is outside bounds, it is restricted to the valid range.
-  private static int getPythonStringIndex(int index, int stringLength) {
+  private static int getClampedIndex(int index, int length) {
     if (index < 0) {
-      index += stringLength;
+      index += length;
     }
-    return Math.max(Math.min(index, stringLength), 0);
+    return Math.max(Math.min(index, length), 0);
   }
 
   // Emulate Python substring function
   // It converts out of range indices, and never fails
   private static String getPythonSubstring(String str, int start, int end) {
-    start = getPythonStringIndex(start, str.length());
-    end = getPythonStringIndex(end, str.length());
+    start = getClampedIndex(start, str.length());
+    end = getClampedIndex(end, str.length());
     if (start > end) {
       return "";
     } else {
@@ -225,7 +225,7 @@ public class MethodLibrary {
             end = Type.INTEGER.convert(args[3], "'rfind' argument");
           }
           int subpos = getPythonSubstring(thiz, start, end).lastIndexOf(sub);
-          start = getPythonStringIndex(start, thiz.length());
+          start = getClampedIndex(start, thiz.length());
           return subpos < 0 ? subpos : subpos + start;
         }
       };
@@ -256,7 +256,7 @@ public class MethodLibrary {
             end = Type.INTEGER.convert(args[3], "'find' argument");
           }
           int subpos = getPythonSubstring(thiz, start, end).indexOf(sub);
-          start = getPythonStringIndex(start, thiz.length());
+          start = getClampedIndex(start, thiz.length());
           return subpos < 0 ? subpos : subpos + start;
         }
       };
@@ -369,17 +369,33 @@ public class MethodLibrary {
         }
       };
 
-  // substring operator
-  @SkylarkBuiltin(name = "$substring", hidden = true,
-      doc = "String[<code>start</code>:<code>end</code>] returns a substring.")
-  private static Function substring = new MixedModeFunction("$substring",
+  // slice operator
+  @SkylarkBuiltin(name = "$slice", hidden = true,
+      doc = "x[<code>start</code>:<code>end</code>] returns a slice or a list slice.")
+  private static Function slice = new MixedModeFunction("$slice",
       ImmutableList.of("this", "start", "end"), 3, false) {
     @Override
-    public Object call(Object[] args, FuncallExpression ast) throws ConversionException {
-      String thiz = Type.STRING.convert(args[0], "substring operand");
-      int left = Type.INTEGER.convert(args[1], "substring operand");
-      int right = Type.INTEGER.convert(args[2], "substring operand");
-      return getPythonSubstring(thiz, left, right);
+    public Object call(Object[] args, FuncallExpression ast, Environment env)
+        throws EvalException, ConversionException {
+      int left = Type.INTEGER.convert(args[1], "start operand");
+      int right = Type.INTEGER.convert(args[2], "end operand");
+
+      // Substring
+      if (args[0] instanceof String) {
+        String thiz = Type.STRING.convert(args[0], "substring operand");
+        return getPythonSubstring(thiz, left, right);
+      }
+
+      // List slice
+      List<Object> list = Type.OBJECT_LIST.convert(args[0], "list operand");
+      left = getClampedIndex(left, list.size());
+      right = getClampedIndex(right, list.size());
+
+      List<Object> result = Lists.newArrayList();
+      for (int i = left; i < right; i++) {
+        result.add(list.get(i));
+      }
+      return convert(result, env, ast.getLocation());
     }
   };
 
@@ -500,6 +516,7 @@ public class MethodLibrary {
     }
   };
 
+  // TODO(bazel-team): Use the same type for both Skylark and BUILD files.
   @SuppressWarnings("unchecked")
   private static Iterable<Object> convert(Collection<?> list, Environment env, Location loc)
       throws EvalException {
@@ -937,10 +954,15 @@ public class MethodLibrary {
       .put(endswith, SkylarkType.BOOL)
       .put(startswith, SkylarkType.BOOL)
       .put(strip, SkylarkType.STRING)
-      .put(substring, SkylarkType.STRING)
+      .put(slice, SkylarkType.STRING)
       .put(count, SkylarkType.INT)
       .build();
 
+  public static final Map<Function, SkylarkType> listPureFunctions = ImmutableMap
+      .<Function, SkylarkType>builder()
+      .put(slice, SkylarkType.LIST)
+      .build();
+
   public static final List<Function> listFunctions = ImmutableList.of(append, extend);
 
   public static final Map<Function, SkylarkType> dictFunctions = ImmutableMap
@@ -986,6 +1008,7 @@ public class MethodLibrary {
     setupMethodEnvironment(env, Map.class, dictFunctions.keySet());
     env.registerFunction(String.class, index.getName(), index);
     setupMethodEnvironment(env, String.class, stringFunctions.keySet());
+    setupMethodEnvironment(env, List.class, listPureFunctions.keySet());
     if (env.isSkylarkEnabled()) {
       env.registerFunction(SkylarkList.class, index.getName(), index);
       setupMethodEnvironment(env, skylarkGlobalFunctions.keySet());
@@ -1033,5 +1056,9 @@ public class MethodLibrary {
     Map<String, SkylarkType> string = new HashMap<>();
     setupValidationEnvironment(stringFunctions, string);
     builtIn.put(SkylarkType.STRING, string);
+
+    Map<String, SkylarkType> list = new HashMap<>();
+    setupValidationEnvironment(listPureFunctions, string);
+    builtIn.put(SkylarkType.LIST, list);
   }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/syntax/FuncallExpression.java b/src/main/java/com/google/devtools/build/lib/syntax/FuncallExpression.java
index c31d4343af..41fe33238f 100644
--- a/src/main/java/com/google/devtools/build/lib/syntax/FuncallExpression.java
+++ b/src/main/java/com/google/devtools/build/lib/syntax/FuncallExpression.java
@@ -247,7 +247,7 @@ public final class FuncallExpression extends Expression {
 
   @Override
   public String toString() {
-    if (func.getName().equals("$substring")) {
+    if (func.getName().equals("$slice")) {
       return obj + "[" + args.get(0) + ":" + args.get(1) + "]";
     }
     if (func.getName().equals("$index")) {
diff --git a/src/main/java/com/google/devtools/build/lib/syntax/Parser.java b/src/main/java/com/google/devtools/build/lib/syntax/Parser.java
index 45eae4c4e0..68a4425ef8 100644
--- a/src/main/java/com/google/devtools/build/lib/syntax/Parser.java
+++ b/src/main/java/com/google/devtools/build/lib/syntax/Parser.java
@@ -697,7 +697,7 @@ class Parser {
       return makeFuncallExpression(receiver, new Ident("$index"), args,
                                    start, token.right);
     }
-    // This is a substring
+    // This is a slice (or substring)
     expect(TokenKind.COLON);
     int loc2 = token.left;
     if (token.kind == TokenKind.RBRACKET) {
@@ -708,7 +708,7 @@ class Parser {
     expect(TokenKind.RBRACKET);
 
     args.add(setLocation(new Argument.Positional(endExpr), loc2, endExpr));
-    return makeFuncallExpression(receiver, new Ident("$substring"), args,
+    return makeFuncallExpression(receiver, new Ident("$slice"), args,
                                  start, token.right);
   }
 
diff --git a/src/test/java/com/google/devtools/build/lib/syntax/ParserTest.java b/src/test/java/com/google/devtools/build/lib/syntax/ParserTest.java
index 0e69effe1e..e08f6d47ee 100644
--- a/src/test/java/com/google/devtools/build/lib/syntax/ParserTest.java
+++ b/src/test/java/com/google/devtools/build/lib/syntax/ParserTest.java
@@ -214,7 +214,7 @@ public class ParserTest extends AbstractParserTestCase {
   @Test
   public void testSubstring() throws Exception {
     FuncallExpression e = (FuncallExpression) parseExpr("'FOO.CC'[:].lower()[1:]");
-    assertEquals("$substring", e.getFunction().getName());
+    assertEquals("$slice", e.getFunction().getName());
     assertThat(e.getArguments()).hasSize(2);
 
     e = (FuncallExpression) parseExpr("'FOO.CC'.lower()[1:].startswith('oo')");
@@ -222,7 +222,7 @@ public class ParserTest extends AbstractParserTestCase {
     assertThat(e.getArguments()).hasSize(1);
 
     e = (FuncallExpression) parseExpr("'FOO.CC'[1:][:2]");
-    assertEquals("$substring", e.getFunction().getName());
+    assertEquals("$slice", e.getFunction().getName());
     assertThat(e.getArguments()).hasSize(2);
   }