Add support for tp_iter and tp_iternext

src/capi/errors.cpp

@@ -133,6 +133,10 @@ extern "C" PyObject* PyErr_SetFromErrno(PyObject* exc) noexcept {
     return PyErr_SetFromErrnoWithFilenameObject(exc, NULL);
 }
 
+extern "C" void PyErr_SetNone(PyObject* exception) noexcept {
+    PyErr_SetObject(exception, (PyObject*)NULL);
+}
+
 /* Call when an exception has occurred but there is no way for Python
    to handle it.  Examples: exception in __del__ or during GC. */
 extern "C" void PyErr_WriteUnraisable(PyObject* obj) noexcept {

src/capi/typeobject.cpp

@@ -79,6 +79,18 @@ RICHCMP_WRAPPER(ne, Py_NE)
 RICHCMP_WRAPPER(gt, Py_GT)
 RICHCMP_WRAPPER(ge, Py_GE)
 
+static PyObject* wrap_next(PyObject* self, PyObject* args, void* wrapped) {
+    unaryfunc func = (unaryfunc)wrapped;
+    PyObject* res;
+
+    if (!check_num_args(args, 0))
+        return NULL;
+    res = (*func)(self);
+    if (res == NULL && !PyErr_Occurred())
+        PyErr_SetNone(PyExc_StopIteration);
+    return res;
+}
+
 static PyObject* wrap_coercefunc(PyObject* self, PyObject* args, void* wrapped) noexcept {
     coercion func = (coercion)wrapped;
     PyObject* other, *res;
@@ -581,7 +593,37 @@ static PyObject* slot_tp_richcompare(PyObject* self, PyObject* other, int op) no
     return Py_NotImplemented;
 }
 
-PyObject* slot_tp_new(PyTypeObject* self, PyObject* args, PyObject* kwds) noexcept {
+static PyObject* slot_tp_iter(PyObject* self) noexcept {
+    PyObject* func, *res;
+    static PyObject* iter_str, *getitem_str;
+
+    func = lookup_method(self, "__iter__", &iter_str);
+    if (func != NULL) {
+        PyObject* args;
+        args = res = PyTuple_New(0);
+        if (args != NULL) {
+            res = PyObject_Call(func, args, NULL);
+            Py_DECREF(args);
+        }
+        Py_DECREF(func);
+        return res;
+    }
+    PyErr_Clear();
+    func = lookup_method(self, "__getitem__", &getitem_str);
+    if (func == NULL) {
+        PyErr_Format(PyExc_TypeError, "'%.200s' object is not iterable", Py_TYPE(self)->tp_name);
+        return NULL;
+    }
+    Py_DECREF(func);
+    return PySeqIter_New(self);
+}
+
+static PyObject* slot_tp_iternext(PyObject* self) noexcept {
+    static PyObject* next_str;
+    return call_method(self, "next", &next_str, "()");
+}
+
+static PyObject* slot_tp_new(PyTypeObject* self, PyObject* args, PyObject* kwds) noexcept {
     try {
         // TODO: runtime ICs?
         Box* new_attr = typeLookup(self, _new_str, NULL);
@@ -1033,6 +1075,9 @@ static slotdef slotdefs[]
         TPSLOT("__gt__", tp_richcompare, slot_tp_richcompare, richcmp_gt, "x.__gt__(y) <==> x>y"),
         TPSLOT("__ge__", tp_richcompare, slot_tp_richcompare, richcmp_ge, "x.__ge__(y) <==> x>=y"),
 
+        TPSLOT("__iter__", tp_iter, slot_tp_iter, wrap_unaryfunc, "x.__iter__() <==> iter(x)"),
+        TPSLOT("next", tp_iternext, slot_tp_iternext, wrap_next, "x.next() -> the next value, or raise StopIteration"),
+
         FLSLOT("__init__", tp_init, slot_tp_init, (wrapperfunc)wrap_init, "x.__init__(...) initializes x; "
                                                                           "see help(type(x)) for signature",
                PyWrapperFlag_KEYWORDS),
@@ -1693,8 +1738,6 @@ extern "C" int PyType_Ready(PyTypeObject* cls) noexcept {
         RELEASE_ASSERT(cls->tp_flags & Py_TPFLAGS_CHECKTYPES, "Pyston doesn't yet support non-checktypes behavior");
     }
 
-    RELEASE_ASSERT(cls->tp_iter == NULL, "");
-    RELEASE_ASSERT(cls->tp_iternext == NULL, "");
     RELEASE_ASSERT(cls->tp_descr_get == NULL, "");
     RELEASE_ASSERT(cls->tp_descr_set == NULL, "");
     RELEASE_ASSERT(cls->tp_free == NULL || cls->tp_free == PyObject_Del || cls->tp_free == PyObject_GC_Del, "");

src/runtime/capi.cpp

@@ -663,14 +663,20 @@ finally:
 }
 
 void checkAndThrowCAPIException() {
-    Box* value = cur_thread_state.curexc_value;
-    if (value) {
+    Box* _type = cur_thread_state.curexc_type;
+    if (!_type)
+        assert(!cur_thread_state.curexc_value);
+
+    if (_type) {
         RELEASE_ASSERT(cur_thread_state.curexc_traceback == NULL, "unsupported");
-        Box* _type = cur_thread_state.curexc_type;
         BoxedClass* type = static_cast<BoxedClass*>(_type);
         assert(isInstance(_type, type_cls) && isSubclass(static_cast<BoxedClass*>(type), BaseException)
                && "Only support throwing subclass of BaseException for now");
 
+        Box* value = cur_thread_state.curexc_value;
+        if (!value)
+            value = None;
+
         // This is similar to PyErr_NormalizeException:
         if (!isInstance(value, type)) {
             if (value->cls == tuple_cls) {

src/runtime/iterobject.cpp

@@ -96,6 +96,15 @@ Box* iterwrapperNext(Box* s) {
     return r;
 }
 
+extern "C" PyObject* PySeqIter_New(PyObject* seq) noexcept {
+    try {
+        return new BoxedSeqIter(seq);
+    } catch (Box* e) {
+        PyErr_SetObject(e->cls, e);
+        return NULL;
+    }
+}
+
 void setupIter() {
     seqiter_cls = new BoxedHeapClass(object_cls, NULL, 0, sizeof(BoxedSeqIter), false);
     seqiter_cls->giveAttr("__name__", boxStrConstant("iterator"));

test/test_extension/slots_test.c

@@ -9,6 +9,14 @@ typedef struct {
     int n;
 } slots_tester_object;
 
+typedef struct {
+    PyObject_HEAD;
+
+    slots_tester_object* obj;
+
+    int m;
+} slots_tester_iterobj;
+
 static PyObject *
 slots_tester_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
 {
@@ -110,6 +118,7 @@ static PySequenceMethods slots_tester_seq_as_sequence = {
     (objobjproc)0,             /* sq_contains */
 };
 
+static slots_tester_iterobj* slots_tester_iter(slots_tester_object *obj);
 
 static PyTypeObject slots_tester_seq = {
     PyVarObject_HEAD_INIT(NULL, 0)
@@ -138,7 +147,7 @@ static PyTypeObject slots_tester_seq = {
     0,                                  /* tp_clear */
     (richcmpfunc)slots_tester_seq_richcmp, /* tp_richcompare */
     0,                                  /* tp_weaklistoffset */
-    0,                                  /* tp_iter */
+    (getiterfunc)slots_tester_iter,                                  /* tp_iter */
     0,                                  /* tp_iternext */
     0,                                  /* tp_methods */
     0,                                  /* tp_members */
@@ -154,6 +163,73 @@ static PyTypeObject slots_tester_seq = {
     0,                                  /* tp_free */
 };
 
+PyObject* slots_testeriter_iternext(slots_tester_iterobj* iter) {
+    iter->m++;
+    if (iter->m < iter->obj->n) {
+        return PyInt_FromLong(iter->m);
+    }
+    return NULL;
+}
+
+static int
+slots_testeriter_init(PyObject *self, PyObject *args, PyObject *kwds)
+{
+    printf("slots_testeriter_seq.__init__, %d\n", ((slots_tester_object*)self)->n);
+
+    return 0;
+}
+
+static PyTypeObject slots_tester_seqiter = {
+    PyVarObject_HEAD_INIT(NULL, 0)
+    "slots_test.slots_tester_seqiter",            /* tp_name */
+    sizeof(slots_tester_iterobj),          /* tp_basicsize */
+    0,                                  /* tp_itemsize */
+    /* methods */
+    0,                                  /* tp_dealloc */
+    0,                                  /* tp_print */
+    0,                                  /* tp_getattr */
+    0,                                  /* tp_setattr */
+    0,                                  /* tp_compare */
+    (reprfunc)slots_tester_seq_repr,        /* tp_repr */
+    0,                                  /* tp_as_number */
+    0,          /* tp_as_sequence */
+    0,                                  /* tp_as_mapping */
+    0, /* tp_hash */
+    0,     /* tp_call */
+    0,     /* tp_str */
+    0,                                  /* tp_getattro */
+    0,                                  /* tp_setattro */
+    0,                                  /* tp_as_buffer */
+    Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE,                 /* tp_flags */
+    0,                   /* tp_doc */
+    0,                                  /* tp_traverse */
+    0,                                  /* tp_clear */
+    0, /* tp_richcompare */
+    0,                                  /* tp_weaklistoffset */
+    0,                                  /* tp_iter */
+    (iternextfunc)slots_testeriter_iternext,                                  /* tp_iternext */
+    0,                                  /* tp_methods */
+    0,                                  /* tp_members */
+    0,                                  /* tp_getset */
+    0,                                  /* tp_base */
+    0,                                  /* tp_dict */
+    0,                                  /* tp_descr_get */
+    0,                                  /* tp_descr_set */
+    0,                                  /* tp_dictoffset */
+    slots_testeriter_init,                  /* tp_init */
+    0,                 /* tp_alloc */
+    0,                   /* tp_new */
+    0,                                  /* tp_free */
+};
+
+static slots_tester_iterobj* slots_tester_iter(slots_tester_object *obj) {
+    slots_tester_iterobj* rtn = PyObject_New(slots_tester_iterobj, &slots_tester_seqiter);
+    Py_INCREF(obj);
+    rtn->obj = obj;
+    rtn->m = 0;
+    return rtn;
+}
+
 static PyTypeObject slots_tester_nonsubclassable = {
     PyVarObject_HEAD_INIT(NULL, 0)
     "slots_test.slots_tester_nonsubclassable",            /* tp_name */
@@ -618,6 +694,10 @@ initslots_test(void)
     if (res < 0)
         return;
 
+    res = PyType_Ready(&slots_tester_seqiter);
+    if (res < 0)
+        return;
+
     res = PyType_Ready(&slots_tester_map);
     if (res < 0)
         return;

test/tests/capi_slots.py

@@ -39,6 +39,9 @@ for i in xrange(3):
     print hex(t)
     print oct(t)
 
+for i in slots_test.SlotsTesterSeq(6):
+    print i
+
 su = slots_test.SlotsTesterSub(5)
 print su