Merge pull request #75 from undingen/strip

Implement str.lstrip and str.rstrip

src/runtime/str.cpp

@@ -380,34 +380,55 @@ Box* strSplit2(BoxedString* self, BoxedString* sep) {
     }
 }
 
-Box* strStrip(BoxedString* self) {
+Box* strStrip1(BoxedString* self) {
     assert(self->cls == str_cls);
+    return new BoxedString(llvm::StringRef(self->s).trim(" \t\n\r\f\v"));
+}
 
-    const std::string& s = self->s;
-    int n = s.size();
-
-    int strip_beginning = 0;
-    while (strip_beginning < n) {
-        char c = s[strip_beginning];
-        if (c == ' ' || c == '\t' || c == '\n' || c == '\r' || c == '\f' || c == '\v')
-            strip_beginning++;
-        else
-            break;
+Box* strStrip2(BoxedString* self, Box* chars) {
+    assert(self->cls == str_cls);
+
+    if (chars->cls == str_cls) {
+        return new BoxedString(llvm::StringRef(self->s).trim(static_cast<BoxedString*>(chars)->s));
+    } else if (chars->cls == none_cls) {
+        return strStrip1(self);
+    } else {
+        raiseExcHelper(TypeError, "strip arg must be None, str or unicode");
     }
+}
 
-    if (strip_beginning == n)
-        return boxStrConstant("");
+Box* strLStrip1(BoxedString* self) {
+    assert(self->cls == str_cls);
+    return new BoxedString(llvm::StringRef(self->s).ltrim(" \t\n\r\f\v"));
+}
 
-    int strip_end = 0;
-    while (strip_end < n) {
-        char c = s[n - strip_end - 1];
-        if (c == ' ' || c == '\t' || c == '\n' || c == '\r' || c == '\f' || c == '\v')
-            strip_end++;
-        else
-            break;
+Box* strLStrip2(BoxedString* self, Box* chars) {
+    assert(self->cls == str_cls);
+
+    if (chars->cls == str_cls) {
+        return new BoxedString(llvm::StringRef(self->s).ltrim(static_cast<BoxedString*>(chars)->s));
+    } else if (chars->cls == none_cls) {
+        return strLStrip1(self);
+    } else {
+        raiseExcHelper(TypeError, "lstrip arg must be None, str or unicode");
     }
+}
+
+Box* strRStrip1(BoxedString* self) {
+    assert(self->cls == str_cls);
+    return new BoxedString(llvm::StringRef(self->s).rtrim(" \t\n\r\f\v"));
+}
+
+Box* strRStrip2(BoxedString* self, Box* chars) {
+    assert(self->cls == str_cls);
 
-    return new BoxedString(s.substr(strip_beginning, n - strip_beginning - strip_end));
+    if (chars->cls == str_cls) {
+        return new BoxedString(llvm::StringRef(self->s).rtrim(static_cast<BoxedString*>(chars)->s));
+    } else if (chars->cls == none_cls) {
+        return strRStrip1(self);
+    } else {
+        raiseExcHelper(TypeError, "rstrip arg must be None, str or unicode");
+    }
 }
 
 Box* strContains(BoxedString* self, Box* elt) {
@@ -541,7 +562,19 @@ void setupStr() {
     str_cls->giveAttr("__nonzero__", new BoxedFunction(boxRTFunction((void*)strNonzero, NULL, 1, false)));
 
     str_cls->giveAttr("lower", new BoxedFunction(boxRTFunction((void*)strLower, STR, 1, false)));
-    str_cls->giveAttr("strip", new BoxedFunction(boxRTFunction((void*)strStrip, STR, 1, false)));
+
+    CLFunction* strStrip = boxRTFunction((void*)strStrip1, STR, 1, false);
+    addRTFunction(strStrip, (void*)strStrip2, STR, 2, false);
+    str_cls->giveAttr("strip", new BoxedFunction(strStrip));
+
+    CLFunction* strLStrip = boxRTFunction((void*)strLStrip1, STR, 1, false);
+    addRTFunction(strLStrip, (void*)strLStrip2, STR, 2, false);
+    str_cls->giveAttr("lstrip", new BoxedFunction(strLStrip));
+
+    CLFunction* strRStrip = boxRTFunction((void*)strRStrip1, STR, 1, false);
+    addRTFunction(strRStrip, (void*)strRStrip2, STR, 2, false);
+    str_cls->giveAttr("rstrip", new BoxedFunction(strRStrip));
+
     str_cls->giveAttr("__contains__", new BoxedFunction(boxRTFunction((void*)strContains, BOXED_BOOL, 2, false)));
 
     str_cls->giveAttr("__add__", new BoxedFunction(boxRTFunction((void*)strAdd, NULL, 2, false)));

test/tests/str_functions.py

@@ -20,7 +20,11 @@ print map(bool, ["hello", "", "world"])
 if "":
     print "bad"
 
-print repr(" \t\n\v\ftest \t\n\v\f".strip())
+testStr = " \t\n\v\ftest \t\n\v\f"
+print repr(testStr.strip()), repr(testStr.lstrip()), repr(testStr.rstrip())
+for c in [None, " ", "\t\f", "test"]:
+    print repr(testStr.strip(c)), repr(testStr.lstrip(c)), repr(testStr.rstrip(c))
+
 
 for pattern in ["hello", "o w", "nope"]:
     print pattern in "hello world"