Implement str.replace

runtime/str.go

@@ -453,6 +453,70 @@ func strNew(f *Frame, t *Type, args Args, _ KWArgs) (*Object, *BaseException) {
 	return s.ToObject(), nil
 }
 
+// strReplace returns a copy of the string s with the first n non-overlapping
+// instances of old replaced by sub. If old is empty, it matches at the
+// beginning of the string. If n < 0, there is no limit on the number of
+// replacements.
+func strReplace(f *Frame, args Args, _ KWArgs) (*Object, *BaseException) {
+	var raised *BaseException
+	// TODO: Support unicode replace.
+	expectedTypes := []*Type{StrType, StrType, StrType, ObjectType}
+	argc := len(args)
+	if argc == 3 {
+		expectedTypes = expectedTypes[:argc]
+	}
+	if raised := checkMethodArgs(f, "replace", args, expectedTypes...); raised != nil {
+		return nil, raised
+	}
+	n := -1
+	if argc == 4 {
+		n, raised = ToIntValue(f, args[3])
+		if raised != nil {
+			return nil, raised
+		}
+	}
+	s := toStrUnsafe(args[0]).Value()
+	// Returns early if no need to replace.
+	if n == 0 {
+		return NewStr(s).ToObject(), nil
+	}
+
+	old := toStrUnsafe(args[1]).Value()
+	sub := toStrUnsafe(args[2]).Value()
+	numBytes := len(s)
+	// Even if s and old is blank, replace should return sub, except n is negative.
+	// This is CPython specific behavior.
+	if numBytes == 0 && old == "" && n >= 0 {
+		return NewStr("").ToObject(), nil
+	}
+	// If old is non-blank, pass to strings.Replace.
+	if len(old) > 0 {
+		return NewStr(strings.Replace(s, old, sub, n)).ToObject(), nil
+	}
+
+	// If old is blank, insert sub after every bytes on s and beginning.
+	if n < 0 {
+		n = numBytes + 1
+	}
+	// Insert sub at beginning.
+	buf := bytes.Buffer{}
+	buf.WriteString(sub)
+	n--
+	// Insert after every byte.
+	i := 0
+	for n > 0 && i < numBytes {
+		buf.WriteByte(s[i])
+		buf.WriteString(sub)
+		i++
+		n--
+	}
+	// Write the remaining string.
+	if i < numBytes {
+		buf.WriteString(s[i:])
+	}
+	return NewStr(buf.String()).ToObject(), nil
+}
+
 func strRepr(_ *Frame, o *Object) (*Object, *BaseException) {
 	s := toStrUnsafe(o).Value()
 	buf := bytes.Buffer{}
@@ -602,6 +666,7 @@ func initStrType(dict map[string]*Object) {
 	dict["split"] = newBuiltinFunction("split", strSplit).ToObject()
 	dict["startswith"] = newBuiltinFunction("startswith", strStartsWith).ToObject()
 	dict["strip"] = newBuiltinFunction("strip", strStrip).ToObject()
+	dict["replace"] = newBuiltinFunction("replace", strReplace).ToObject()
 	dict["rstrip"] = newBuiltinFunction("rstrip", strRStrip).ToObject()
 	dict["title"] = newBuiltinFunction("title", strTitle).ToObject()
 	dict["upper"] = newBuiltinFunction("upper", strUpper).ToObject()

runtime/str_test.go

@@ -258,6 +258,16 @@ func TestStrMethods(t *testing.T) {
 			return NewLong(big.NewInt(2)).ToObject(), nil
 		}).ToObject(),
 	}))
+	intIntType := newTestClass("IntInt", []*Type{ObjectType}, newStringDict(map[string]*Object{
+		"__int__": newBuiltinFunction("__int__", func(f *Frame, _ Args, _ KWArgs) (*Object, *BaseException) {
+			return NewInt(2).ToObject(), nil
+		}).ToObject(),
+	}))
+	longIntType := newTestClass("LongInt", []*Type{ObjectType}, newStringDict(map[string]*Object{
+		"__int__": newBuiltinFunction("__int__", func(f *Frame, _ Args, _ KWArgs) (*Object, *BaseException) {
+			return NewLong(big.NewInt(2)).ToObject(), nil
+		}).ToObject(),
+	}))
 	cases := []struct {
 		methodName string
 		args       Args
@@ -362,6 +372,41 @@ func TestStrMethods(t *testing.T) {
 		{"strip", wrapArgs("foo", "bar", "baz"), nil, mustCreateException(TypeErrorType, "'strip' of 'str' requires 2 arguments")},
 		{"strip", wrapArgs("\xfboo", NewUnicode("o")), nil, mustCreateException(UnicodeDecodeErrorType, "'utf8' codec can't decode byte 0xfb in position 0")},
 		{"strip", wrapArgs("foo", NewUnicode("o")), NewUnicode("f").ToObject(), nil},
+		{"replace", wrapArgs("one!two!three!", "!", "@", 1), NewStr("one@two!three!").ToObject(), nil},
+		{"replace", wrapArgs("one!two!three!", "!", ""), NewStr("onetwothree").ToObject(), nil},
+		{"replace", wrapArgs("one!two!three!", "!", "@", 2), NewStr("one@two@three!").ToObject(), nil},
+		{"replace", wrapArgs("one!two!three!", "!", "@", 3), NewStr("one@two@three@").ToObject(), nil},
+		{"replace", wrapArgs("one!two!three!", "!", "@", 4), NewStr("one@two@three@").ToObject(), nil},
+		{"replace", wrapArgs("one!two!three!", "!", "@", 0), NewStr("one!two!three!").ToObject(), nil},
+		{"replace", wrapArgs("one!two!three!", "!", "@"), NewStr("one@two@three@").ToObject(), nil},
+		{"replace", wrapArgs("one!two!three!", "x", "@"), NewStr("one!two!three!").ToObject(), nil},
+		{"replace", wrapArgs("one!two!three!", "x", "@", 2), NewStr("one!two!three!").ToObject(), nil},
+		{"replace", wrapArgs("\xd0\xb2\xd0\xbe\xd0\xbb", "", "\x00", -1), NewStr("\x00\xd0\x00\xb2\x00\xd0\x00\xbe\x00\xd0\x00\xbb\x00").ToObject(), nil},
+		{"replace", wrapArgs("\xd0\xb2\xd0\xbe\xd0\xbb", "", "\x01\x02", -1), NewStr("\x01\x02\xd0\x01\x02\xb2\x01\x02\xd0\x01\x02\xbe\x01\x02\xd0\x01\x02\xbb\x01\x02").ToObject(), nil},
+		{"replace", wrapArgs("abc", "", "-"), NewStr("-a-b-c-").ToObject(), nil},
+		{"replace", wrapArgs("abc", "", "-", 3), NewStr("-a-b-c").ToObject(), nil},
+		{"replace", wrapArgs("abc", "", "-", 0), NewStr("abc").ToObject(), nil},
+		{"replace", wrapArgs("", "", ""), NewStr("").ToObject(), nil},
+		{"replace", wrapArgs("", "", "a"), NewStr("a").ToObject(), nil},
+		{"replace", wrapArgs("abc", "a", "--", 0), NewStr("abc").ToObject(), nil},
+		{"replace", wrapArgs("abc", "xy", "--"), NewStr("abc").ToObject(), nil},
+		{"replace", wrapArgs("123", "123", ""), NewStr("").ToObject(), nil},
+		{"replace", wrapArgs("123123", "123", ""), NewStr("").ToObject(), nil},
+		{"replace", wrapArgs("123x123", "123", ""), NewStr("x").ToObject(), nil},
+		{"replace", wrapArgs("one!two!three!", "!", "@", NewLong(big.NewInt(1))), NewStr("one@two!three!").ToObject(), nil},
+		{"replace", wrapArgs("foobar", "bar", "baz", newObject(intIntType)), NewStr("foobaz").ToObject(), nil},
+		{"replace", wrapArgs("foobar", "bar", "baz", newObject(longIntType)), NewStr("foobaz").ToObject(), nil},
+		{"replace", wrapArgs("", "", "x"), NewStr("x").ToObject(), nil},
+		{"replace", wrapArgs("", "", "x", -1), NewStr("x").ToObject(), nil},
+		{"replace", wrapArgs("", "", "x", 0), NewStr("").ToObject(), nil},
+		{"replace", wrapArgs("", "", "x", 1), NewStr("").ToObject(), nil},
+		{"replace", wrapArgs("", "", "x", 1000), NewStr("").ToObject(), nil},
+		// TODO: Support unicode substring.
+		{"replace", wrapArgs("foobar", "", NewUnicode("bar")), nil, mustCreateException(TypeErrorType, "'replace' requires a 'str' object but received a 'unicode'")},
+		{"replace", wrapArgs("foobar", NewUnicode("bar"), ""), nil, mustCreateException(TypeErrorType, "'replace' requires a 'str' object but received a 'unicode'")},
+		{"replace", wrapArgs("foobar", "bar", "baz", None), nil, mustCreateException(TypeErrorType, "an integer is required")},
+		{"replace", wrapArgs("foobar", "bar", "baz", newObject(intIndexType)), nil, mustCreateException(TypeErrorType, "an integer is required")},
+		{"replace", wrapArgs("foobar", "bar", "baz", newObject(longIndexType)), nil, mustCreateException(TypeErrorType, "an integer is required")},
 		{"rstrip", wrapArgs("foo "), NewStr("foo").ToObject(), nil},
 		{"rstrip", wrapArgs(" foo bar "), NewStr(" foo bar").ToObject(), nil},
 		{"rstrip", wrapArgs("foo foo", "o"), NewStr("foo f").ToObject(), nil},

testing/str_test.py

@@ -105,6 +105,85 @@ def __index__(self):
 assert "%x" % 0x1f == "1f"
 assert "%X" % 0xffff == "FFFF"
 
+# Test replace
+assert 'one!two!three!'.replace('!', '@', 1) == 'one@two!three!'
+assert 'one!two!three!'.replace('!', '') == 'onetwothree'
+assert 'one!two!three!'.replace('!', '@', 2) == 'one@two@three!'
+assert 'one!two!three!'.replace('!', '@', 3) == 'one@two@three@'
+assert 'one!two!three!'.replace('!', '@', 4) == 'one@two@three@'
+assert 'one!two!three!'.replace('!', '@', 0) == 'one!two!three!'
+assert 'one!two!three!'.replace('!', '@') == 'one@two@three@'
+assert 'one!two!three!'.replace('x', '@') == 'one!two!three!'
+assert 'one!two!three!'.replace('x', '@', 2) == 'one!two!three!'
+assert 'abc'.replace('', '-') == '-a-b-c-'
+assert 'abc'.replace('', '-', 3) == '-a-b-c'
+assert 'abc'.replace('', '-', 0) == 'abc'
+assert ''.replace('', '') == ''
+assert ''.replace('', 'a') == 'a'
+assert 'abc'.replace('a', '--', 0) == 'abc'
+assert 'abc'.replace('xy', '--') == 'abc'
+assert '123'.replace('123', '') == ''
+assert '123123'.replace('123', '') == ''
+assert '123x123'.replace('123', '') == 'x'
+assert "\xd0\xb2\xd0\xbe\xd0\xbb".replace('', '\0') == "\x00\xd0\x00\xb2\x00\xd0\x00\xbe\x00\xd0\x00\xbb\x00"
+assert "\xd0\xb2\xd0\xbe\xd0\xbb".replace('', '\1\2') == '\x01\x02\xd0\x01\x02\xb2\x01\x02\xd0\x01\x02\xbe\x01\x02\xd0\x01\x02\xbb\x01\x02'
+
+class S(str):
+  pass
+
+s = S('abc')
+assert type(s.replace(s, s)) is str
+assert type(s.replace('x', 'y')) is str
+assert type(s.replace('x', 'y', 0)) is str
+# CPython only, pypy supposed to be same as Go
+assert ''.replace('', 'x') == 'x'
+assert ''.replace('', 'x', -1) == 'x'
+assert ''.replace('', 'x', 0) == ''
+assert ''.replace('', 'x', 1) == ''
+assert ''.replace('', 'x', 1000) == ''
+try:
+  ''.replace(None, '')
+  raise AssertionError
+except TypeError:
+  pass
+try:
+  ''.replace('', None)
+  raise AssertionError
+except TypeError:
+  pass
+try:
+  ''.replace('', '', None)
+  raise AssertionError
+except TypeError:
+  pass
+
+class A(object):
+  def __int__(self):
+    return 3
+class AL(object):
+  def __int__(self):
+    return 3L
+
+class B(object):
+  def __index__(self):
+    return 3
+class BL(object):
+  def __index__(self):
+    return 3L
+
+assert 'aaaaa'.replace('a', 'b', A()) == 'bbbaa'
+assert 'aaaaa'.replace('a', 'b', AL()) == 'bbbaa'
+try:
+  'aaaaa'.replace('a', 'b', B())
+  raise AssertionError
+except TypeError:
+  pass
+try:
+  'aaaaa'.replace('a', 'b', BL())
+  raise AssertionError
+except TypeError:
+  pass
+
 # Test zfill
 assert '123'.zfill(2) == '123'
 assert '123'.zfill(3) == '123'