Implement str.zfill (#124)

runtime/str.go

@@ -561,6 +561,7 @@ func initStrType(dict map[string]*Object) {
 	dict["rstrip"] = newBuiltinFunction("rstrip", strRStrip).ToObject()
 	dict["title"] = newBuiltinFunction("title", strTitle).ToObject()
 	dict["upper"] = newBuiltinFunction("upper", strUpper).ToObject()
+	dict["zfill"] = newBuiltinFunction("zfill", strZFill).ToObject()
 	StrType.slots.Add = &binaryOpSlot{strAdd}
 	StrType.slots.Contains = &binaryOpSlot{strContains}
 	StrType.slots.Eq = &binaryOpSlot{strEq}
@@ -796,6 +797,32 @@ func strUpper(f *Frame, args Args, kwargs KWArgs) (*Object, *BaseException) {
 	return NewStr(strings.ToUpper(s)).ToObject(), nil
 }
 
+func strZFill(f *Frame, args Args, _ KWArgs) (*Object, *BaseException) {
+	if raised := checkMethodArgs(f, "zfill", args, StrType, ObjectType); raised != nil {
+		return nil, raised
+	}
+	s := toStrUnsafe(args[0]).Value()
+	l := len(s)
+	o, raised := IntType.Call(f, Args{args[1]}, nil)
+	if raised != nil {
+		return nil, raised
+	}
+	width := toIntUnsafe(o).Value()
+	if width <= l {
+		return args[0], nil
+	}
+	buf := bytes.Buffer{}
+	buf.Grow(width)
+	if l > 0 && (s[0] == '-' || s[0] == '+') {
+		buf.WriteByte(s[0])
+		s = s[1:]
+		width--
+	}
+	buf.WriteString(strings.Repeat("0", width-len(s)))
+	buf.WriteString(s)
+	return NewStr(buf.String()).ToObject(), nil
+}
+
 func init() {
 	InternStr("")
 	for i := 0; i < 256; i++ {

runtime/str_test.go

@@ -346,6 +346,28 @@ func TestStrMethods(t *testing.T) {
 		{"upper", wrapArgs("aBC"), NewStr("ABC").ToObject(), nil},
 		{"upper", wrapArgs("abc def", 123), nil, mustCreateException(TypeErrorType, "'upper' of 'str' requires 1 arguments")},
 		{"upper", wrapArgs(123), nil, mustCreateException(TypeErrorType, "unbound method upper() must be called with str instance as first argument (got int instance instead)")},
+		{"zfill", wrapArgs("123", 2), NewStr("123").ToObject(), nil},
+		{"zfill", wrapArgs("123", 3), NewStr("123").ToObject(), nil},
+		{"zfill", wrapArgs("123", 4), NewStr("0123").ToObject(), nil},
+		{"zfill", wrapArgs("+123", 3), NewStr("+123").ToObject(), nil},
+		{"zfill", wrapArgs("+123", 4), NewStr("+123").ToObject(), nil},
+		{"zfill", wrapArgs("+123", 5), NewStr("+0123").ToObject(), nil},
+		{"zfill", wrapArgs("-123", 3), NewStr("-123").ToObject(), nil},
+		{"zfill", wrapArgs("-123", 4), NewStr("-123").ToObject(), nil},
+		{"zfill", wrapArgs("-123", 5), NewStr("-0123").ToObject(), nil},
+		{"zfill", wrapArgs("123", NewLong(big.NewInt(3))), NewStr("123").ToObject(), nil},
+		{"zfill", wrapArgs("123", NewLong(big.NewInt(5))), NewStr("00123").ToObject(), nil},
+		{"zfill", wrapArgs("", 0), NewStr("").ToObject(), nil},
+		{"zfill", wrapArgs("", 1), NewStr("0").ToObject(), nil},
+		{"zfill", wrapArgs("", 3), NewStr("000").ToObject(), nil},
+		{"zfill", wrapArgs("", -1), NewStr("").ToObject(), nil},
+		{"zfill", wrapArgs("34", 1), NewStr("34").ToObject(), nil},
+		{"zfill", wrapArgs("34", 4), NewStr("0034").ToObject(), nil},
+		{"zfill", wrapArgs("34", None), nil, mustCreateException(TypeErrorType, "int() argument must be a string or a number, not 'NoneType'")},
+		{"zfill", wrapArgs("", True), NewStr("0").ToObject(), nil},
+		{"zfill", wrapArgs("", False), NewStr("").ToObject(), nil},
+		{"zfill", wrapArgs("34", NewStr("test")), nil, mustCreateException(ValueErrorType, "invalid literal for int() with base 10: test")},
+		{"zfill", wrapArgs("34"), nil, mustCreateException(TypeErrorType, "'zfill' of 'str' requires 2 arguments")},
 	}
 	for _, cas := range cases {
 		testCase := invokeTestCase{args: cas.args, want: cas.want, wantExc: cas.wantExc}

testing/str_test.py

@@ -28,3 +28,29 @@ assert "%%" % tuple() == "%"
 assert "%r" % "abc" == "'abc'"
 assert "%x" % 0x1f == "1f"
 assert "%X" % 0xffff == "FFFF"
+
+# Test zfill
+assert '123'.zfill(2) == '123'
+assert '123'.zfill(3) == '123'
+assert '123'.zfill(4) == '0123'
+assert '+123'.zfill(3) == '+123'
+assert '+123'.zfill(4) == '+123'
+assert '+123'.zfill(5) == '+0123'
+assert '-123'.zfill(3) == '-123'
+assert '-123'.zfill(4) == '-123'
+assert '-123'.zfill(5) == '-0123'
+assert ''.zfill(3) == '000'
+assert '34'.zfill(1) == '34'
+assert '34'.zfill(4) == '0034'
+
+try:
+  '123'.zfill()
+  raise AssertionError
+except TypeError:
+  pass
+
+class A(object):
+  def __int__(self):
+    return 3
+
+assert '3'.zfill(A()) == '003'