Support two- and three-arg startswith/endswith

src/core/cfg.cpp

@@ -19,8 +19,6 @@
 #include <cstdio>
 #include <cstdlib>
 
-#include "llvm/Support/ErrorHandling.h" // For llvm_unreachable
-
 #include "analysis/scoping_analysis.h"
 #include "core/ast.h"
 #include "core/options.h"

src/runtime/capi.cpp

@@ -19,6 +19,8 @@
 
 #include "Python.h"
 
+#include "llvm/Support/ErrorHandling.h" // For llvm_unreachable
+
 #include "capi/types.h"
 #include "core/threading.h"
 #include "core/types.h"
@@ -675,6 +677,11 @@ void setCAPIException(const ExcInfo& e) {
     cur_thread_state.curexc_traceback = e.traceback;
 }
 
+void throwCAPIException() {
+    checkAndThrowCAPIException();
+    llvm_unreachable("No exception was thrown?");
+}
+
 void checkAndThrowCAPIException() {
     Box* _type = cur_thread_state.curexc_type;
     if (!_type)
@@ -1377,6 +1384,29 @@ extern "C" PyObject* PyCFunction_NewEx(PyMethodDef* ml, PyObject* self, PyObject
     Py_FatalError("unimplemented");
 }
 
+extern "C" int _PyEval_SliceIndex(PyObject* v, Py_ssize_t* pi) noexcept {
+    if (v != NULL) {
+        Py_ssize_t x;
+        if (PyInt_Check(v)) {
+            /* XXX(nnorwitz): I think PyInt_AS_LONG is correct,
+               however, it looks like it should be AsSsize_t.
+               There should be a comment here explaining why.
+            */
+            x = PyInt_AS_LONG(v);
+        } else if (PyIndex_Check(v)) {
+            x = PyNumber_AsSsize_t(v, NULL);
+            if (x == -1 && PyErr_Occurred())
+                return 0;
+        } else {
+            PyErr_SetString(PyExc_TypeError, "slice indices must be integers or "
+                                             "None or have an __index__ method");
+            return 0;
+        }
+        *pi = x;
+    }
+    return 1;
+}
+
 BoxedModule* importTestExtension(const std::string& name) {
     std::string pathname_name = "test/test_extension/" + name + ".pyston.so";
     const char* pathname = pathname_name.c_str();

src/runtime/capi.h

@@ -23,6 +23,7 @@ class BoxedModule;
 BoxedModule* importTestExtension(const std::string&);
 
 void checkAndThrowCAPIException();
+void throwCAPIException() __attribute__((noreturn));
 struct ExcInfo;
 void setCAPIException(const ExcInfo& e);
 }

src/runtime/str.cpp

@@ -1494,7 +1494,9 @@ Box* strContains(BoxedString* self, Box* elt) {
     return True;
 }
 
-Box* strStartswith(BoxedString* self, Box* elt) {
+Box* strStartswith(BoxedString* self, Box* elt, Box* start, Box** _args) {
+    Box* end = _args[0];
+
     if (self->cls != str_cls)
         raiseExcHelper(TypeError, "descriptor 'startswith' requires a 'str' object but received a '%s'",
                        getTypeName(self)->c_str());
@@ -1502,12 +1504,44 @@ Box* strStartswith(BoxedString* self, Box* elt) {
     if (elt->cls != str_cls)
         raiseExcHelper(TypeError, "expected a character buffer object");
 
+    Py_ssize_t istart = 0, iend = PY_SSIZE_T_MAX;
+    if (start) {
+        int r = _PyEval_SliceIndex(start, &istart);
+        if (!r)
+            throwCAPIException();
+    }
+
+    if (end) {
+        int r = _PyEval_SliceIndex(end, &iend);
+        if (!r)
+            throwCAPIException();
+    }
+
     BoxedString* sub = static_cast<BoxedString*>(elt);
 
-    return boxBool(startswith(self->s, sub->s));
+    Py_ssize_t n = self->s.size();
+    iend = std::min(iend, n);
+    if (iend < 0)
+        iend += n;
+    if (iend < 0)
+        iend = 0;
+
+    if (istart < 0)
+        istart += n;
+    if (istart < 0)
+        istart = 0;
+
+    Py_ssize_t compare_len = iend - istart;
+    if (compare_len < 0)
+        return False;
+    if (sub->s.size() > compare_len)
+        return False;
+    return boxBool(self->s.compare(istart, sub->s.size(), sub->s) == 0);
 }
 
-Box* strEndswith(BoxedString* self, Box* elt) {
+Box* strEndswith(BoxedString* self, Box* elt, Box* start, Box** _args) {
+    Box* end = _args[0];
+
     if (self->cls != str_cls)
         raiseExcHelper(TypeError, "descriptor 'endswith' requires a 'str' object but received a '%s'",
                        getTypeName(self)->c_str());
@@ -1515,9 +1549,41 @@ Box* strEndswith(BoxedString* self, Box* elt) {
     if (elt->cls != str_cls)
         raiseExcHelper(TypeError, "expected a character buffer object");
 
+    Py_ssize_t istart = 0, iend = PY_SSIZE_T_MAX;
+    if (start) {
+        int r = _PyEval_SliceIndex(start, &istart);
+        if (!r)
+            throwCAPIException();
+    }
+
+    if (end) {
+        int r = _PyEval_SliceIndex(end, &iend);
+        if (!r)
+            throwCAPIException();
+    }
+
     BoxedString* sub = static_cast<BoxedString*>(elt);
 
-    return boxBool(endswith(self->s, sub->s));
+    Py_ssize_t n = self->s.size();
+    iend = std::min(iend, n);
+    if (iend < 0)
+        iend += n;
+    if (iend < 0)
+        iend = 0;
+
+    if (istart < 0)
+        istart += n;
+    if (istart < 0)
+        istart = 0;
+
+    Py_ssize_t compare_len = iend - istart;
+    if (compare_len < 0)
+        return False;
+    if (sub->s.size() > compare_len)
+        return False;
+    // XXX: this line is the only difference between startswith and endswith:
+    istart += compare_len - sub->s.size();
+    return boxBool(self->s.compare(istart, sub->s.size(), sub->s) == 0);
 }
 
 Box* strFind(BoxedString* self, Box* elt, Box* _start) {
@@ -1790,8 +1856,10 @@ void setupStr() {
 
     str_cls->giveAttr("__contains__", new BoxedFunction(boxRTFunction((void*)strContains, BOXED_BOOL, 2)));
 
-    str_cls->giveAttr("startswith", new BoxedFunction(boxRTFunction((void*)strStartswith, BOXED_BOOL, 2)));
-    str_cls->giveAttr("endswith", new BoxedFunction(boxRTFunction((void*)strEndswith, BOXED_BOOL, 2)));
+    str_cls->giveAttr("startswith",
+                      new BoxedFunction(boxRTFunction((void*)strStartswith, BOXED_BOOL, 4, 2, 0, 0), { NULL, NULL }));
+    str_cls->giveAttr("endswith",
+                      new BoxedFunction(boxRTFunction((void*)strEndswith, BOXED_BOOL, 4, 2, 0, 0), { NULL, NULL }));
 
     str_cls->giveAttr("find",
                       new BoxedFunction(boxRTFunction((void*)strFind, BOXED_INT, 3, 1, false, false), { boxInt(0) }));

test/tests/str_functions.py

@@ -81,3 +81,8 @@ print "hello world"[False:True:True]
 
 print "{hello}".format(hello="world")
 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)