[GR-11836] Implementation of str.zfill is missing.

PullRequest: graalpython/203

Author: ppisl

graalpython/com.oracle.graal.python.test/src/tests/test_string.py

@@ -475,7 +475,8 @@ class subtype(self.__class__.type2test):
             else:
                 obj = subtype(obj)
                 realresult = getattr(obj, methodname)(*args)
-                self.assertIsNot(obj, realresult)
+                self.assertTrue(obj is not realresult)
+                #self.assertIsNot(obj, realresult)
 
     # check that obj.method(*args) raises exc
     def checkraises(self, exc, obj, methodname, *args):
@@ -688,6 +689,32 @@ def test_isprintable(self):
         self.assertTrue('\U0001F46F'.isprintable())
         # self.assertFalse('\U000E0020'.isprintable())
 
+    def test_zfill(self):
+        self.checkequal('123', '123', 'zfill', 2)
+        self.checkequal('123', '123', 'zfill', 3)
+        self.checkequal('0123', '123', 'zfill', 4)
+        self.checkequal('+123', '+123', 'zfill', 3)
+        self.checkequal('+123', '+123', 'zfill', 4)
+        self.checkequal('+0123', '+123', 'zfill', 5)
+        self.checkequal('-123', '-123', 'zfill', 3)
+        self.checkequal('-123', '-123', 'zfill', 4)
+        self.checkequal('-0123', '-123', 'zfill', 5)
+        self.checkequal('000', '', 'zfill', 3)
+        self.checkequal('34', '34', 'zfill', 1)
+        self.checkequal('0034', '34', 'zfill', 4)
+
+        self.checkraises(TypeError, '123', 'zfill')
+        
+    def test_zfill_specialization(self):
+        self.checkequal('123', '123', 'zfill', True)
+
+        class MyIndexable(object):
+            def __init__(self, value):
+                self.value = value
+            def __index__(self):
+                return self.value
+
+        self.checkequal('0123', '123', 'zfill', MyIndexable(4))
 
 def test_same_id():
     empty_ids = set([id(str()) for i in range(100)])

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/str/StringBuiltins.java

@@ -86,6 +86,8 @@
 import com.oracle.graal.python.nodes.function.builtins.PythonUnaryBuiltinNode;
 import com.oracle.graal.python.nodes.object.GetClassNode;
 import com.oracle.graal.python.nodes.truffle.PythonArithmeticTypes;
+import com.oracle.graal.python.nodes.util.CastToIndexNode;
+import com.oracle.graal.python.nodes.util.CastToIntegerFromIndexNode;
 import com.oracle.graal.python.runtime.exception.PException;
 import com.oracle.graal.python.runtime.formatting.StringFormatter;
 import com.oracle.truffle.api.CompilerDirectives;
@@ -1543,6 +1545,61 @@ boolean doString(String self) {
         }
     }
 
+    @Builtin(name = "zfill", fixedNumOfPositionalArgs = 2)
+    @GenerateNodeFactory
+    @TypeSystemReference(PythonArithmeticTypes.class)
+    abstract static class ZFillNode extends PythonBinaryBuiltinNode {
+
+        public abstract String executeObject(String self, Object x);
+
+        @Specialization
+        public String doString(String self, long width) {
+            return zfill(self, (int) width);
+        }
+
+        @Specialization
+        public String doString(String self, PInt width,
+                        @Cached("create()") CastToIndexNode toIndexNode) {
+            return zfill(self, toIndexNode.execute(width));
+        }
+
+        @Specialization
+        public String doString(String self, Object width,
+                        @Cached("create()") CastToIntegerFromIndexNode widthCast,
+                        @Cached("create()") ZFillNode recursiveNode) {
+            return recursiveNode.executeObject(self, widthCast.execute(width));
+        }
+
+        private static String zfill(String self, int width) {
+            int len = self.length();
+            if (len >= width) {
+                return self;
+            }
+            char[] chars = new char[width];
+            int nzeros = width - len;
+            int i = 0;
+            int sStart = 0;
+            if (len > 0) {
+                char start = self.charAt(0);
+                if (start == '+' || start == '-') {
+                    chars[0] = start;
+                    i += 1;
+                    nzeros++;
+                    sStart = 1;
+                }
+            }
+            for (; i < nzeros; i++) {
+                chars[i] = '0';
+            }
+            self.getChars(sStart, len, chars, i);
+            return new String(chars);
+        }
+
+        public static ZFillNode create() {
+            return StringBuiltinsFactory.ZFillNodeFactory.create();
+        }
+    }
+
     @Builtin(name = __GETITEM__, fixedNumOfPositionalArgs = 2)
     @GenerateNodeFactory
     @TypeSystemReference(PythonArithmeticTypes.class)