Substring function (#180)

* Substring function * substring -> substr
d5 · Apr 9, 2019 · 09009d4 · 09009d4
1 parent 1a0c06b
commit 09009d4
Show file tree

Hide file tree

Showing 2 changed files with 83 additions and 0 deletions.
diff --git a/stdlib/text.go b/stdlib/text.go
@@ -32,6 +32,7 @@ var textModule = map[string]objects.Object{
 	"last_index_any": &objects.UserFunction{Name: "last_index_any", Value: FuncASSRI(strings.LastIndexAny)}, // last_index_any(s, chars) => int
 	"repeat":         &objects.UserFunction{Name: "repeat", Value: textRepeat},                              // repeat(s, count) => string
 	"replace":        &objects.UserFunction{Name: "replace", Value: textReplace},                            // replace(s, old, new, n) => string
+	"substr":         &objects.UserFunction{Name: "substr", Value: textSubstring},                           // substring(s, lower, upper) => string
 	"split":          &objects.UserFunction{Name: "split", Value: FuncASSRSs(strings.Split)},                // split(s, sep) => [string]
 	"split_after":    &objects.UserFunction{Name: "split_after", Value: FuncASSRSs(strings.SplitAfter)},     // split_after(s, sep) => [string]
 	"split_after_n":  &objects.UserFunction{Name: "split_after_n", Value: FuncASSIRSs(strings.SplitAfterN)}, // split_after_n(s, sep, n) => [string]
@@ -376,6 +377,69 @@ func textReplace(args ...objects.Object) (ret objects.Object, err error) {
 	return
 }
 
+func textSubstring(args ...objects.Object) (ret objects.Object, err error) {
+	argslen := len(args)
+	if argslen != 2 && argslen != 3 {
+		err = objects.ErrWrongNumArguments
+		return
+	}
+
+	s1, ok := objects.ToString(args[0])
+	if !ok {
+		err = objects.ErrInvalidArgumentType{
+			Name:     "first",
+			Expected: "string(compatible)",
+			Found:    args[0].TypeName(),
+		}
+		return
+	}
+
+	i2, ok := objects.ToInt(args[1])
+	if !ok {
+		err = objects.ErrInvalidArgumentType{
+			Name:     "second",
+			Expected: "int(compatible)",
+			Found:    args[1].TypeName(),
+		}
+		return
+	}
+
+	strlen := len(s1)
+	i3 := strlen
+	if argslen == 3 {
+		i3, ok = objects.ToInt(args[2])
+		if !ok {
+			err = objects.ErrInvalidArgumentType{
+				Name:     "third",
+				Expected: "int(compatible)",
+				Found:    args[2].TypeName(),
+			}
+			return
+		}
+	}
+
+	if i2 > i3 {
+		err = objects.ErrInvalidIndexType
+		return
+	}
+
+	if i2 < 0 {
+		i2 = 0
+	} else if i2 > strlen {
+		i2 = strlen
+	}
+
+	if i3 < 0 {
+		i3 = 0
+	} else if i3 > strlen {
+		i3 = strlen
+	}
+
+	ret = &objects.String{Value: s1[i2:i3]}
+
+	return
+}
+
 func textRepeat(args ...objects.Object) (ret objects.Object, err error) {
 	if len(args) != 2 {
 		return nil, objects.ErrWrongNumArguments

diff --git a/stdlib/text_test.go b/stdlib/text_test.go
@@ -244,3 +244,22 @@ func TestTextRepeat(t *testing.T) {
 	module(t, "text").call("repeat", "1", "12").expect("111111111111")
 	module(t, "text").call("repeat", "1", "13").expectError()
 }
+
+func TestSubstr(t *testing.T) {
+	module(t, "text").call("substr", "", 0, 0).expect("")
+	module(t, "text").call("substr", "abcdef", 0, 3).expect("abc")
+	module(t, "text").call("substr", "abcdef", 0, 6).expect("abcdef")
+	module(t, "text").call("substr", "abcdef", 0, 10).expect("abcdef")
+	module(t, "text").call("substr", "abcdef", -10, 10).expect("abcdef")
+	module(t, "text").call("substr", "abcdef", 0).expect("abcdef")
+	module(t, "text").call("substr", "abcdef", 3).expect("def")
+
+	module(t, "text").call("substr", "", 10, 0).expectError()
+	module(t, "text").call("substr", "", "10", 0).expectError()
+	module(t, "text").call("substr", "", 10, "0").expectError()
+	module(t, "text").call("substr", "", "10", "0").expectError()
+
+	module(t, "text").call("substr", 0, 0, 1).expect("0")
+	module(t, "text").call("substr", 123, 0, 1).expect("1")
+	module(t, "text").call("substr", 123.456, 4, 7).expect("456")
+}