Add a few missing str, list and set methods

from_cpython/Objects/stringobject.c

@@ -23,6 +23,41 @@ PyObject * _do_string_format(PyObject *self, PyObject *args, PyObject *kwargs) {
     return do_string_format(self, args, kwargs);
 }
 
+PyObject *
+string_count(PyStringObject *self, PyObject *args)
+{
+    PyObject *sub_obj;
+    const char *str = PyString_AS_STRING(self), *sub;
+    Py_ssize_t sub_len;
+    Py_ssize_t start = 0, end = PY_SSIZE_T_MAX;
+
+    if (!stringlib_parse_args_finds("count", args, &sub_obj, &start, &end))
+        return NULL;
+
+    if (PyString_Check(sub_obj)) {
+        sub = PyString_AS_STRING(sub_obj);
+        sub_len = PyString_GET_SIZE(sub_obj);
+    }
+#ifdef Py_USING_UNICODE
+    else if (PyUnicode_Check(sub_obj)) {
+        Py_ssize_t count;
+        count = PyUnicode_Count((PyObject *)self, sub_obj, start, end);
+        if (count == -1)
+            return NULL;
+        else
+            return PyInt_FromSsize_t(count);
+    }
+#endif
+    else if (PyObject_AsCharBuffer(sub_obj, &sub, &sub_len))
+        return NULL;
+
+    ADJUST_INDICES(start, end, PyString_GET_SIZE(self));
+
+    return PyInt_FromSsize_t(
+        stringlib_count(str + start, end - start, sub, sub_len, PY_SSIZE_T_MAX)
+        );
+}
+
 PyObject * string_split(PyStringObject *self, PyObject *args)
 {
     Py_ssize_t len = PyString_GET_SIZE(self), n;
@@ -130,6 +165,20 @@ string_find(PyStringObject *self, PyObject *args)
     return PyInt_FromSsize_t(result);
 }
 
+PyObject *
+string_index(PyStringObject *self, PyObject *args)
+{
+    Py_ssize_t result = string_find_internal(self, args, +1);
+    if (result == -2)
+        return NULL;
+    if (result == -1) {
+        PyErr_SetString(PyExc_ValueError,
+                        "substring not found");
+        return NULL;
+    }
+    return PyInt_FromSsize_t(result);
+}
+
 PyObject*
 string_splitlines(PyStringObject *self, PyObject *args)
 {

src/runtime/list.cpp

@@ -589,12 +589,29 @@ Box* listCount(BoxedList* self, Box* elt) {
     return boxInt(count);
 }
 
-Box* listIndex(BoxedList* self, Box* elt) {
+Box* listIndex(BoxedList* self, Box* elt, BoxedInt* _start, Box** args) {
     LOCK_REGION(self->lock.asRead());
 
-    int size = self->size;
+    BoxedInt* _stop = (BoxedInt*)args[0];
+    RELEASE_ASSERT(!_start || _start->cls == int_cls, "");
+    RELEASE_ASSERT(!_stop || _stop->cls == int_cls, "");
 
-    for (int i = 0; i < size; i++) {
+    int64_t start = _start ? _start->n : 0;
+    int64_t stop = _stop ? _stop->n : self->size;
+
+    if (start < 0) {
+        start += self->size;
+        if (start < 0)
+            start = 0;
+    }
+
+    if (stop < 0) {
+        stop += self->size;
+        if (stop < 0)
+            stop = 0;
+    }
+
+    for (int64_t i = start; i < stop; i++) {
         Box* e = self->elts->elts[i];
         Box* cmp = compareInternal(e, elt, AST_TYPE::Eq, NULL);
         bool b = nonzero(cmp);
@@ -840,7 +857,8 @@ void setupList() {
                        new BoxedFunction(boxRTFunction((void*)listInit, UNKNOWN, 2, 1, false, false), { None }));
 
     list_cls->giveAttr("count", new BoxedFunction(boxRTFunction((void*)listCount, BOXED_INT, 2)));
-    list_cls->giveAttr("index", new BoxedFunction(boxRTFunction((void*)listIndex, BOXED_INT, 2)));
+    list_cls->giveAttr(
+        "index", new BoxedFunction(boxRTFunction((void*)listIndex, BOXED_INT, 4, 2, false, false), { NULL, NULL }));
     list_cls->giveAttr("remove", new BoxedFunction(boxRTFunction((void*)listRemove, NONE, 2)));
     list_cls->giveAttr("reverse", new BoxedFunction(boxRTFunction((void*)listReverse, NONE, 1)));
     list_cls->freeze();

src/runtime/set.cpp

@@ -252,6 +252,29 @@ Box* setUnion(BoxedSet* self, BoxedTuple* args) {
     return rtn;
 }
 
+static BoxedSet* setIntersection2(BoxedSet* self, Box* container) {
+    assert(self->cls == set_cls);
+
+    BoxedSet* rtn = new BoxedSet();
+    for (auto elt : container->pyElements()) {
+        if (self->s.count(elt))
+            rtn->s.insert(elt);
+    }
+    return rtn;
+}
+
+Box* setIntersection(BoxedSet* self, BoxedTuple* args) {
+    if (!isSubclass(self->cls, set_cls))
+        raiseExcHelper(TypeError, "descriptor 'intersection' requires a 'set' object but received a '%s'",
+                       getTypeName(self));
+
+    BoxedSet* rtn = self;
+    for (auto container : args->pyElements()) {
+        rtn = setIntersection2(rtn, container);
+    }
+    return rtn;
+}
+
 Box* setCopy(BoxedSet* self) {
     assert(self->cls == set_cls);
 
@@ -269,6 +292,18 @@ Box* setNonzero(BoxedSet* self) {
     return boxBool(self->s.size());
 }
 
+Box* setHash(BoxedSet* self) {
+    RELEASE_ASSERT(isSubclass(self->cls, frozenset_cls), "");
+
+    int64_t rtn = 1927868237L;
+    for (Box* e : self->s) {
+        BoxedInt* h = hash(e);
+        assert(isSubclass(h->cls, int_cls));
+        rtn ^= h->n + 0x9e3779b9 + (rtn << 6) + (rtn >> 2);
+    }
+
+    return boxInt(rtn);
+}
 
 } // namespace set
 
@@ -338,12 +373,16 @@ void setupSet() {
     set_cls->giveAttr("__nonzero__", new BoxedFunction(boxRTFunction((void*)setNonzero, BOXED_BOOL, 1)));
     frozenset_cls->giveAttr("__nonzero__", set_cls->getattr("__nonzero__"));
 
+    frozenset_cls->giveAttr("__hash__", new BoxedFunction(boxRTFunction((void*)setHash, BOXED_INT, 1)));
+
     set_cls->giveAttr("add", new BoxedFunction(boxRTFunction((void*)setAdd, NONE, 2)));
     set_cls->giveAttr("remove", new BoxedFunction(boxRTFunction((void*)setRemove, NONE, 2)));
 
     set_cls->giveAttr("clear", new BoxedFunction(boxRTFunction((void*)setClear, NONE, 1)));
     set_cls->giveAttr("update", new BoxedFunction(boxRTFunction((void*)setUpdate, NONE, 1, 0, true, false)));
     set_cls->giveAttr("union", new BoxedFunction(boxRTFunction((void*)setUnion, UNKNOWN, 1, 0, true, false)));
+    set_cls->giveAttr("intersection",
+                      new BoxedFunction(boxRTFunction((void*)setIntersection, UNKNOWN, 1, 0, true, false)));
 
     set_cls->giveAttr("copy", new BoxedFunction(boxRTFunction((void*)setCopy, UNKNOWN, 1)));
 

src/runtime/str.cpp

@@ -35,9 +35,11 @@
 #include "runtime/types.h"
 #include "runtime/util.h"
 
+extern "C" PyObject* string_count(PyStringObject* self, PyObject* args) noexcept;
 extern "C" PyObject* string_split(PyStringObject* self, PyObject* args) noexcept;
 extern "C" PyObject* string_rsplit(PyStringObject* self, PyObject* args) noexcept;
 extern "C" PyObject* string_find(PyStringObject* self, PyObject* args) noexcept;
+extern "C" PyObject* string_index(PyStringObject* self, PyObject* args) noexcept;
 extern "C" PyObject* string_rfind(PyStringObject* self, PyObject* args) noexcept;
 extern "C" PyObject* string_splitlines(PyStringObject* self, PyObject* args) noexcept;
 
@@ -2080,8 +2082,8 @@ Box* strEncode(BoxedString* self, Box* encoding, Box* error) {
     if (error_str && error_str->cls != str_cls)
         raiseExcHelper(TypeError, "encode() argument 2 must be string, not '%s'", getTypeName(error_str));
 
-    Box* result
-        = PyCodec_Encode(self, encoding_str ? encoding_str->s.c_str() : NULL, error_str ? error_str->s.c_str() : NULL);
+    Box* result = PyCodec_Encode(self, encoding_str ? encoding_str->s.c_str() : PyUnicode_GetDefaultEncoding(),
+                                 error_str ? error_str->s.c_str() : NULL);
     checkAndThrowCAPIException();
     return result;
 }
@@ -2160,48 +2162,6 @@ Box* strIter(BoxedString* self) {
     return new BoxedStringIterator(self);
 }
 
-int64_t strCount2Unboxed(BoxedString* self, Box* elt) {
-    assert(isSubclass(self->cls, str_cls));
-
-    if (elt->cls != str_cls)
-        raiseExcHelper(TypeError, "expected a character buffer object");
-
-    const std::string& s = self->s;
-    const std::string& pattern = static_cast<BoxedString*>(elt)->s;
-
-    int found = 0;
-    size_t start = 0;
-    while (start < s.size()) {
-        size_t next = s.find(pattern, start);
-        if (next == std::string::npos)
-            break;
-
-        found++;
-        start = next + pattern.size();
-    }
-    return found;
-}
-
-Box* strCount2(BoxedString* self, Box* elt) {
-    return boxInt(strCount2Unboxed(self, elt));
-}
-
-Box* strIndex(BoxedString* self, Box* elt) {
-    assert(isSubclass(self->cls, str_cls));
-
-    if (elt->cls != str_cls)
-        raiseExcHelper(TypeError, "expected a character buffer object");
-
-    const std::string& s = self->s;
-    const std::string& pattern = static_cast<BoxedString*>(elt)->s;
-
-    size_t idx = s.find(pattern, 0);
-    if (idx == std::string::npos)
-        raiseExcHelper(ValueError, "substring not found");
-
-    return boxInt(idx);
-}
-
 extern "C" PyObject* PyString_FromString(const char* s) noexcept {
     return boxStrConstant(s);
 }
@@ -2434,9 +2394,11 @@ void strDestructor(Box* b) {
 }
 
 static PyMethodDef string_methods[] = {
+    { "count", (PyCFunction)string_count, METH_VARARGS, NULL },
     { "split", (PyCFunction)string_split, METH_VARARGS, NULL },
     { "rsplit", (PyCFunction)string_rsplit, METH_VARARGS, NULL },
     { "find", (PyCFunction)string_find, METH_VARARGS, NULL },
+    { "index", (PyCFunction)string_index, METH_VARARGS, NULL },
     { "rfind", (PyCFunction)string_rfind, METH_VARARGS, NULL },
     { "expandtabs", (PyCFunction)string_expandtabs, METH_VARARGS, NULL },
     { "splitlines", (PyCFunction)string_splitlines, METH_VARARGS, NULL },
@@ -2534,11 +2496,6 @@ void setupStr() {
         str_cls->giveAttr(md.ml_name, new BoxedMethodDescriptor(&md, str_cls));
     }
 
-    CLFunction* count = boxRTFunction((void*)strCount2Unboxed, INT, 2);
-    addRTFunction(count, (void*)strCount2, BOXED_INT);
-    str_cls->giveAttr("count", new BoxedFunction(count));
-    str_cls->giveAttr("index", new BoxedFunction(boxRTFunction((void*)strIndex, BOXED_INT, 2)));
-
     str_cls->giveAttr("__new__", new BoxedFunction(boxRTFunction((void*)strNew, UNKNOWN, 2, 1, false, false),
                                                    { boxStrConstant("") }));
 

test/tests/list.py

@@ -53,6 +53,15 @@ for i in xrange(5):
 list_index = [1, 2, 3, 4, 5]
 for i in xrange(1, 6):
     assert list_index.index(i) == i-1
+    try:
+        print list_index.index(i, 3, 4)
+    except ValueError as e:
+        print e
+    try:
+        print list_index.index(i, -1, -1)
+    except ValueError as e:
+        print e
+        
 assert list_index.index(3) == 2
 assert [1, '2'].index('2') == 1
 

test/tests/set.py

@@ -80,8 +80,8 @@ s2 = s.copy()
 s.add(1)
 print s, s2
 
-s1 = set([3])
-s2 = set([1])
-print sorted(s1.union(s2))
-print sorted(s1.union(range(5, 7)))
-print s2.union([], [], [], [])
+s1 = set([3, 5])
+s2 = set([1, 5])
+print sorted(s1.union(s2)), sorted(s1.intersection(s2))
+print sorted(s1.union(range(5, 7))), sorted(s1.intersection(range(5, 7)))
+print sorted(s2.union([], [], [], [])), sorted(s2.intersection())

test/tests/str_functions.py

@@ -38,6 +38,8 @@ for c in "hello world":
 
 for c in "hello world":
     print c, "hello world".count(c)
+    print c, "hello world".count(c, 1, 2)
+    print c, "hello world".count(c, 2, 5)
 
 for i in xrange(1, 10):
     for j in xrange(1, 4):
@@ -132,6 +134,8 @@ for c in "hello world":
     gc.collect()
 
 print "hello world".index("world")
+print "hello world".index("world", 1, 30)
+print "hello world".index("l", 3)
 try:
     print "hello world".index("goodbye")
 except: