From 1ae7af7fbf53fd37bf4a6451c84264b79836ad99 Mon Sep 17 00:00:00 2001
From: Kyle Ellrott <kellrott@gmail.com>
Date: Tue, 10 Jan 2023 15:56:19 -0800
Subject: [PATCH] Adding strip, rstrip and lstrip methods to string class

---
 py/string.go       | 60 ++++++++++++++++++++++++++++++++++++++++++++++
 py/tests/string.py | 11 +++++++++
 2 files changed, 71 insertions(+)

diff --git a/py/string.go b/py/string.go
index d1c0c6f1..a28e6e74 100644
--- a/py/string.go
+++ b/py/string.go
@@ -206,6 +206,18 @@ replaced.`)
 		return Bool(false), nil
 	}, 0, "startswith(prefix[, start[, end]]) -> bool")
 
+	StringType.Dict["strip"] = MustNewMethod("strip", func(self Object, args Tuple, kwargs StringDict) (Object, error) {
+		return self.(String).Strip(args)
+	}, 0, "strip(chars) -> replace chars from begining and end of string")
+
+	StringType.Dict["rstrip"] = MustNewMethod("rstrip", func(self Object, args Tuple, kwargs StringDict) (Object, error) {
+		return self.(String).RStrip(args)
+	}, 0, "rstrip(chars) -> replace chars from end of string")
+
+	StringType.Dict["lstrip"] = MustNewMethod("lstrip", func(self Object, args Tuple, kwargs StringDict) (Object, error) {
+		return self.(String).LStrip(args)
+	}, 0, "lstrip(chars) -> replace chars from begining of string")
+
 }
 
 // Type of this object
@@ -679,6 +691,54 @@ func (s String) Replace(args Tuple) (Object, error) {
 	return String(strings.Replace(string(s), old, new, cnt)), nil
 }
 
+func stripFunc(args Tuple) (func(rune) bool, error) {
+	var (
+		pyval Object = None
+	)
+	err := ParseTuple(args, "|s", &pyval)
+	if err != nil {
+		return nil, err
+	}
+	f := unicode.IsSpace
+	switch v := pyval.(type) {
+	case String:
+		chars := []rune(string(v))
+		f = func(s rune) bool {
+			for _, i := range chars {
+				if s == i {
+					return true
+				}
+			}
+			return false
+		}
+	}
+	return f, nil
+}
+
+func (s String) Strip(args Tuple) (Object, error) {
+	f, err := stripFunc(args)
+	if err != nil {
+		return nil, err
+	}
+	return String(strings.TrimFunc(string(s), f)), nil
+}
+
+func (s String) LStrip(args Tuple) (Object, error) {
+	f, err := stripFunc(args)
+	if err != nil {
+		return nil, err
+	}
+	return String(strings.TrimLeftFunc(string(s), f)), nil
+}
+
+func (s String) RStrip(args Tuple) (Object, error) {
+	f, err := stripFunc(args)
+	if err != nil {
+		return nil, err
+	}
+	return String(strings.TrimRightFunc(string(s), f)), nil
+}
+
 // Check stringerface is satisfied
 var (
 	_ richComparison     = String("")
diff --git a/py/tests/string.py b/py/tests/string.py
index 43487328..8af36ca6 100644
--- a/py/tests/string.py
+++ b/py/tests/string.py
@@ -886,6 +886,17 @@ def index(s, i):
 assert uni[7:7:2] == ''
 assert uni[7:7:3] == ''
 
+doc="string strip methods"
+a = "  adfasd   "
+assert a.rstrip() == "  adfasd"
+assert a.lstrip() == "adfasd   "
+assert a.strip() == "adfasd"
+
+a = " a bada a"
+assert a.rstrip("a ") == " a bad"
+assert a.lstrip("a ") == "bada a"
+assert a.strip("a ") == "bad"
+
 class Index:
     def __index__(self):
         return 1