Add removeprefix/removesuffix to Starlark strings

Implements bazelbuild/starlark#185 in Java Starlark. Closes bazelbuild#14824. PiperOrigin-RevId: 430556229
databricks · Feb 23, 2022 · 0b09e9e · 0b09e9e
1 parent a20797e
commit 0b09e9e
Show file tree

Hide file tree

Showing 2 changed files with 64 additions and 0 deletions.
diff --git a/src/main/java/net/starlark/java/eval/StringModule.java b/src/main/java/net/starlark/java/eval/StringModule.java
@@ -1023,4 +1023,36 @@ public boolean startsWith(String self, Object sub, Object start, Object end)
   private static boolean substringStartsWith(String str, int start, int end, String prefix) {
     return start + prefix.length() <= end && str.startsWith(prefix, start);
   }
+
+  @StarlarkMethod(
+      name = "removeprefix",
+      doc =
+          "If the string starts with <code>prefix</code>, returns a new string with the prefix "
+              + "removed. Otherwise, returns the string.",
+      parameters = {
+        @Param(name = "self", doc = "This string."),
+        @Param(name = "prefix", doc = "The prefix to remove if present."),
+      })
+  public String removePrefix(String self, String prefix) {
+    if (self.startsWith(prefix)) {
+      return self.substring(prefix.length());
+    }
+    return self;
+  }
+
+  @StarlarkMethod(
+      name = "removesuffix",
+      doc =
+          "If the string ends with <code>suffix</code>, returns a new string with the suffix "
+              + "removed. Otherwise, returns the string.",
+      parameters = {
+        @Param(name = "self", doc = "This string."),
+        @Param(name = "suffix", doc = "The suffix to remove if present."),
+      })
+  public String removeSuffix(String self, String suffix) {
+    if (self.endsWith(suffix)) {
+      return self.substring(0, self.length() - suffix.length());
+    }
+    return self;
+  }
 }
diff --git a/src/test/java/net/starlark/java/eval/testdata/string_misc.star b/src/test/java/net/starlark/java/eval/testdata/string_misc.star
@@ -177,3 +177,35 @@ assert_eq("abc" * 0, "")
 assert_eq("abc" * -1, "")
 assert_fails(lambda: "abc" * (1 << 35), "got 34359738368 for repeat, want value in signed 32-bit range")
 assert_fails(lambda: "abc" * (1 << 30), "excessive repeat \\(3 \\* 1073741824 characters\\)")
+
+# removeprefix
+assert_eq("Apricot".removeprefix("Apr"), "icot")
+assert_eq("Apricot".removeprefix("apr"), "Apricot")
+assert_eq("Apricot".removeprefix("A"), "pricot")
+assert_eq("a".removeprefix(""), "a")
+assert_eq("".removeprefix(""), "")
+assert_eq("".removeprefix("a"), "")
+assert_eq("Apricot".removeprefix("pr"), "Apricot")
+assert_eq("AprApricot".removeprefix("Apr"), "Apricot")
+def removeprefix_self_unmodified():
+    original_string = "Apricot"
+    assert_eq(original_string.removeprefix("Apr"), "icot")
+    assert_eq(original_string, "Apricot")
+removeprefix_self_unmodified()
+assert_fails(lambda: "1234".removeprefix(1), "got value of type 'int', want 'string")
+
+# removesuffix
+assert_eq("Apricot".removesuffix("cot"), "Apri")
+assert_eq("Apricot".removesuffix("Cot"), "Apricot")
+assert_eq("Apricot".removesuffix("t"), "Aprico")
+assert_eq("a".removesuffix(""), "a")
+assert_eq("".removesuffix(""), "")
+assert_eq("".removesuffix("a"), "")
+assert_eq("Apricot".removesuffix("co"), "Apricot")
+assert_eq("Apricotcot".removesuffix("cot"), "Apricot")
+def removesuffix_self_unmodified():
+    original_string = "Apricot"
+    assert_eq(original_string.removesuffix("cot"), "Apri")
+    assert_eq(original_string, "Apricot")
+removesuffix_self_unmodified()
+assert_fails(lambda: "1234".removesuffix(4), "got value of type 'int', want 'string")