Merge pull request #267 from tjhance/str-just-funcs

Str just funcs

src/runtime/str.cpp

@@ -976,6 +976,59 @@ extern "C" Box* strNe(BoxedString* lhs, Box* rhs) {
     return boxBool(lhs->s != srhs->s);
 }
 
+#define JUST_LEFT 0
+#define JUST_RIGHT 1
+#define JUST_CENTER 2
+static Box* pad(BoxedString* self, Box* width, Box* fillchar, int justType) {
+    assert(width->cls == int_cls);
+    assert(fillchar->cls == str_cls);
+    assert(static_cast<BoxedString*>(fillchar)->s.size() == 1);
+    int64_t curWidth = self->s.size();
+    int64_t targetWidth = static_cast<BoxedInt*>(width)->n;
+
+    if (curWidth >= targetWidth) {
+        if (self->cls == str_cls) {
+            return self;
+        } else {
+            // If self isn't a string but a subclass of str, then make a new string to return
+            return new BoxedString(self->s);
+        }
+    }
+
+    char c = static_cast<BoxedString*>(fillchar)->s[0];
+
+    int padLeft, padRight;
+    int nNeeded = targetWidth - curWidth;
+    switch (justType) {
+        case JUST_LEFT:
+            padLeft = 0;
+            padRight = nNeeded;
+            break;
+        case JUST_RIGHT:
+            padLeft = nNeeded;
+            padRight = 0;
+            break;
+        case JUST_CENTER:
+            padLeft = nNeeded / 2 + (nNeeded & targetWidth & 1);
+            padRight = nNeeded - padLeft;
+            break;
+    }
+
+    // TODO this is probably slow
+    std::string res = std::string(padLeft, c) + self->s + std::string(padRight, c);
+
+    return new BoxedString(std::move(res));
+}
+extern "C" Box* strLjust(BoxedString* lhs, Box* width, Box* fillchar) {
+    return pad(lhs, width, fillchar, JUST_LEFT);
+}
+extern "C" Box* strRjust(BoxedString* lhs, Box* width, Box* fillchar) {
+    return pad(lhs, width, fillchar, JUST_RIGHT);
+}
+extern "C" Box* strCenter(BoxedString* lhs, Box* width, Box* fillchar) {
+    return pad(lhs, width, fillchar, JUST_CENTER);
+}
+
 extern "C" Box* strLen(BoxedString* self) {
     assert(self->cls == str_cls);
 
@@ -1882,6 +1935,14 @@ void setupStr() {
     str_cls->giveAttr("__eq__", new BoxedFunction(boxRTFunction((void*)strEq, UNKNOWN, 2)));
     str_cls->giveAttr("__ne__", new BoxedFunction(boxRTFunction((void*)strNe, UNKNOWN, 2)));
 
+    BoxedString* spaceChar = new BoxedString(" ");
+    str_cls->giveAttr("ljust",
+                      new BoxedFunction(boxRTFunction((void*)strLjust, UNKNOWN, 3, 1, false, false), { spaceChar }));
+    str_cls->giveAttr("rjust",
+                      new BoxedFunction(boxRTFunction((void*)strRjust, UNKNOWN, 3, 1, false, false), { spaceChar }));
+    str_cls->giveAttr("center",
+                      new BoxedFunction(boxRTFunction((void*)strCenter, UNKNOWN, 3, 1, false, false), { spaceChar }));
+
     str_cls->giveAttr("__getitem__", new BoxedFunction(boxRTFunction((void*)strGetitem, STR, 2)));
 
     str_cls->giveAttr("__iter__", new BoxedFunction(boxRTFunction((void*)strIter, typeFromClass(str_iterator_cls), 1)));

test/tests/str_functions.py

@@ -86,3 +86,33 @@ print "%.3s" % "hello world"
 for i in xrange(-5, 15):
     for j in xrange(-5, 15):
         print i, j, "banana".startswith("ana", i, j), "banana".endswith("ana", i, j)
+
+def test_just_funcs(s, w):
+    t1 = s.ljust(w, 'x')
+    t2 = s.rjust(w, 'x')
+    t3 = s.center(w, 'x')
+
+    t4 = s.ljust(w)
+    t5 = s.rjust(w)
+    t6 = s.center(w)
+
+    print t1, t1 == s, t1 is s, type(t1)
+    print t2, t2 == s, t2 is s, type(t2)
+    print t3, t3 == s, t3 is s, type(t3)
+
+    print t4, t4 == s, t4 is s, type(t4)
+    print t5, t5 == s, t5 is s, type(t5)
+    print t6, t6 == s, t6 is s, type(t6)
+
+test_just_funcs("abcd", 3)
+test_just_funcs("abcd", 4)
+test_just_funcs("abcd", 5)
+test_just_funcs("abcd", 6)
+test_just_funcs("abcd", 7)
+test_just_funcs("abcd", 8)
+test_just_funcs("abcde", 3)
+test_just_funcs("abcde", 4)
+test_just_funcs("abcde", 5)
+test_just_funcs("abcde", 6)
+test_just_funcs("abcde", 7)
+test_just_funcs("abcde", 8)