Merge pull request #41 from undingen/strSplit2

Implement string.split(sep)

src/runtime/str.cpp

@@ -28,6 +28,9 @@
 #include "runtime/types.h"
 #include "runtime/util.h"
 
+#include <llvm/ADT/StringRef.h>
+#include <llvm/ADT/SmallVector.h>
+
 namespace pyston {
 
 extern "C" BoxedString* strAdd(BoxedString* lhs, Box* _rhs) {
@@ -345,6 +348,31 @@ Box* strSplit1(BoxedString* self) {
     return rtn;
 }
 
+Box* strSplit2(BoxedString* self, BoxedString* sep) {
+    assert(self->cls == str_cls);
+
+    if (sep->cls == str_cls) {
+        if (!sep->s.empty()) {
+            llvm::SmallVector<llvm::StringRef, 16> parts;
+            llvm::StringRef(self->s).split(parts, sep->s);
+
+            BoxedList* rtn = new BoxedList();
+            for (const auto &s : parts)
+                listAppendInternal(rtn, boxString(s.str()));
+            return rtn;
+        } else {
+            fprintf(stderr, "ValueError: empty separator\n");
+            raiseExc();
+        }
+    } else if (sep->cls == none_cls) {
+        return strSplit1(self);
+    } else {
+        fprintf(stderr, "TypeError: expected a character buffer object\n");
+        raiseExc();
+    }
+}
+
+
 extern "C" Box* strGetitem(BoxedString* self, Box* slice) {
     if (slice->cls == int_cls) {
         BoxedInt* islice = static_cast<BoxedInt*>(slice);
@@ -390,7 +418,10 @@ void setupStr() {
     str_cls->giveAttr("__getitem__", new BoxedFunction(boxRTFunction((void*)strGetitem, NULL, 2, false)));
 
     str_cls->giveAttr("join", new BoxedFunction(boxRTFunction((void*)strJoin, NULL, 2, false)));
-    str_cls->giveAttr("split", new BoxedFunction(boxRTFunction((void*)strSplit1, LIST, 1, false)));
+
+    CLFunction *strSplit = boxRTFunction((void*)strSplit1, LIST, 1, false);
+    addRTFunction(strSplit, (void*)strSplit2, LIST, 2, false);
+    str_cls->giveAttr("split", new BoxedFunction(strSplit));
 
     CLFunction *__new__ = boxRTFunction((void*)strNew1, NULL, 1, false);
     addRTFunction(__new__, (void*)strNew2, NULL, 2, false);

test/tests/str_functions.py

@@ -7,4 +7,6 @@ print repr("'\"")
 
 print "hello world\tmore\nwords\va\fb\ao".split()
 print "  test  ".split()
-# print "  test  ".split(' ')
+print "  test  ".split(' ')
+print "  test  ".split(None)
+print "1<>2<>3".split('<>')