py: add String.find

sbinet · sbinet · commit ab6c445f6337 · 2022-05-24T19:05:17.000+02:00
Updates #139. Signed-off-by: Sebastien Binet <binet@cern.ch>
diff --git a/py/string.go b/py/string.go
@@ -122,6 +122,43 @@ func fieldsN(s string, n int) []string {
 }
 
 func init() {
+	StringType.Dict["endswith"] = MustNewMethod("endswith", func(self Object, args Tuple) (Object, error) {
+		selfStr := string(self.(String))
+		suffix := []string{}
+		if len(args) > 0 {
+			if s, ok := args[0].(String); ok {
+				suffix = append(suffix, string(s))
+			} else if s, ok := args[0].(Tuple); ok {
+				for _, t := range s {
+					if v, ok := t.(String); ok {
+						suffix = append(suffix, string(v))
+					}
+				}
+			} else {
+				return nil, ExceptionNewf(TypeError, "endswith first arg must be str, unicode, or tuple, not %s", args[0].Type())
+			}
+		} else {
+			return nil, ExceptionNewf(TypeError, "endswith() takes at least 1 argument (0 given)")
+		}
+		for _, s := range suffix {
+			if strings.HasSuffix(selfStr, s) {
+				return Bool(true), nil
+			}
+		}
+		return Bool(false), nil
+	}, 0, "endswith(suffix[, start[, end]]) -> bool")
+
+	StringType.Dict["find"] = MustNewMethod("find", func(self Object, args Tuple) (Object, error) {
+		return self.(String).find(args)
+	}, 0, `find(...)
+S.find(sub[, start[, end]]) -> int
+
+Return the lowest index in S where substring sub is found,
+such that sub is contained within S[start:end].  Optional
+arguments start and end are interpreted as in slice notation.
+
+Return -1 on failure.`)
+
 	StringType.Dict["replace"] = MustNewMethod("replace", func(self Object, args Tuple) (Object, error) {
 		return self.(String).Replace(args)
 	}, 0, `replace(self, old, new, count=-1) -> return a copy with all occurrences of substring old replaced by new.
@@ -169,32 +206,6 @@ replaced.`)
 		return Bool(false), nil
 	}, 0, "startswith(prefix[, start[, end]]) -> bool")
 
-	StringType.Dict["endswith"] = MustNewMethod("endswith", func(self Object, args Tuple) (Object, error) {
-		selfStr := string(self.(String))
-		suffix := []string{}
-		if len(args) > 0 {
-			if s, ok := args[0].(String); ok {
-				suffix = append(suffix, string(s))
-			} else if s, ok := args[0].(Tuple); ok {
-				for _, t := range s {
-					if v, ok := t.(String); ok {
-						suffix = append(suffix, string(v))
-					}
-				}
-			} else {
-				return nil, ExceptionNewf(TypeError, "endswith first arg must be str, unicode, or tuple, not %s", args[0].Type())
-			}
-		} else {
-			return nil, ExceptionNewf(TypeError, "endswith() takes at least 1 argument (0 given)")
-		}
-		for _, s := range suffix {
-			if strings.HasSuffix(selfStr, s) {
-				return Bool(true), nil
-			}
-		}
-		return Bool(false), nil
-	}, 0, "endswith(suffix[, start[, end]]) -> bool")
-
 }
 
 // Type of this object
@@ -578,6 +589,32 @@ func (s String) M__contains__(item Object) (Object, error) {
 	return NewBool(strings.Contains(string(s), string(needle))), nil
 }
 
+func (s String) find(args Tuple) (Object, error) {
+	var (
+		pysub Object
+		pybeg Object = Int(0)
+		pyend Object = Int(len(s))
+		pyfmt        = "s|ii:find"
+	)
+	err := ParseTuple(args, pyfmt, &pysub, &pybeg, &pyend)
+	if err != nil {
+		return nil, err
+	}
+
+	var (
+		beg = int(pybeg.(Int))
+		end = int(pyend.(Int))
+		off = s.slice(0, beg, s.len()).len()
+		str = string(s.slice(beg, end, s.len()))
+		sub = string(pysub.(String))
+		idx = strings.Index(str, sub)
+	)
+	if idx < 0 {
+		return Int(idx), nil
+	}
+	return Int(off + String(str[:idx]).len()), nil
+}
+
 func (s String) Split(args Tuple, kwargs StringDict) (Object, error) {
 	var (
 		pyval Object = None
diff --git a/py/string_test.go b/py/string_test.go
@@ -0,0 +1,79 @@
+// Copyright 2022 The go-python Authors.  All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package py
+
+import "testing"
+
+func TestStringFind(t *testing.T) {
+	for _, tc := range []struct {
+		str string
+		sub string
+		beg int
+		end int
+		idx int
+	}{
+		{
+			str: "hello world",
+			sub: "world",
+			idx: 6,
+		},
+		{
+			str: "hello world",
+			sub: "o",
+			idx: 4,
+		},
+		{
+			str: "hello world",
+			sub: "o",
+			beg: 5,
+			idx: 7,
+		},
+		{
+			str: "hello world",
+			sub: "bye",
+			idx: -1,
+		},
+		{
+			str: "Hello, 世界",
+			sub: "界",
+			idx: 8,
+		},
+		{
+			str: "01234 6789",
+			sub: " ",
+			beg: 6,
+			idx: -1,
+		},
+		{
+			str: "0123456789",
+			sub: "6",
+			beg: 1,
+			end: 6,
+			idx: -1,
+		},
+		{
+			str: "0123456789",
+			sub: "6",
+			beg: 1,
+			end: 7,
+			idx: 6,
+		},
+	} {
+		t.Run(tc.str+":"+tc.sub, func(t *testing.T) {
+			beg := tc.beg
+			end := tc.end
+			if end == 0 {
+				end = len(tc.str)
+			}
+			idx, err := String(tc.str).find(Tuple{String(tc.sub), Int(beg), Int(end)})
+			if err != nil {
+				t.Fatalf("invalid: %+v", err)
+			}
+			if got, want := int(idx.(Int)), tc.idx; got != want {
+				t.Fatalf("got=%d, want=%d", got, want)
+			}
+		})
+	}
+}
