Add mapping methods and corresponding test

src/capi/typeobject.cpp

@@ -97,46 +97,48 @@ static PyObject* wrap_sq_item(PyObject* self, PyObject* args, void* wrapped) noe
     return NULL;
 }
 
-PyObject* slot_tp_new(PyTypeObject* self, PyObject* args, PyObject* kwds) noexcept {
-    try {
-        // TODO: runtime ICs?
-        Box* new_attr = typeLookup(self, _new_str, NULL);
-        assert(new_attr);
-        new_attr = processDescriptor(new_attr, None, self);
+static PyObject* wrap_lenfunc(PyObject* self, PyObject* args, void* wrapped) {
+    lenfunc func = (lenfunc)wrapped;
+    Py_ssize_t res;
 
-        return runtimeCall(new_attr, ArgPassSpec(1, 0, true, true), self, args, kwds, NULL, NULL);
-    } catch (Box* e) {
-        abort();
-    }
+    if (!check_num_args(args, 0))
+        return NULL;
+    res = (*func)(self);
+    if (res == -1 && PyErr_Occurred())
+        return NULL;
+    return PyInt_FromLong((long)res);
 }
 
-PyObject* slot_tp_call(PyObject* self, PyObject* args, PyObject* kwds) noexcept {
-    try {
-        Py_FatalError("this function is untested");
+static PyObject* wrap_objobjargproc(PyObject* self, PyObject* args, void* wrapped) {
+    objobjargproc func = (objobjargproc)wrapped;
+    int res;
+    PyObject* key, *value;
 
-        // TODO: runtime ICs?
-        return runtimeCall(self, ArgPassSpec(0, 0, true, true), args, kwds, NULL, NULL, NULL);
-    } catch (Box* e) {
-        abort();
-    }
+    if (!PyArg_UnpackTuple(args, "", 2, 2, &key, &value))
+        return NULL;
+    res = (*func)(self, key, value);
+    if (res == -1 && PyErr_Occurred())
+        return NULL;
+    Py_INCREF(Py_None);
+    return Py_None;
 }
 
-PyObject* slot_tp_repr(PyObject* self) noexcept {
-    try {
-        return repr(self);
-    } catch (Box* e) {
-        abort();
-    }
-}
+static PyObject* wrap_delitem(PyObject* self, PyObject* args, void* wrapped) {
+    objobjargproc func = (objobjargproc)wrapped;
+    int res;
+    PyObject* key;
 
-PyObject* slot_sq_item(PyObject* self, Py_ssize_t i) noexcept {
-    try {
-        return getitem(self, boxInt(i));
-    } catch (Box* e) {
-        abort();
-    }
+    if (!check_num_args(args, 1))
+        return NULL;
+    key = PyTuple_GET_ITEM(args, 0);
+    res = (*func)(self, key, NULL);
+    if (res == -1 && PyErr_Occurred())
+        return NULL;
+    Py_INCREF(Py_None);
+    return Py_None;
 }
 
+
 static PyObject* lookup_maybe(PyObject* self, const char* attrstr, PyObject** attrobj) {
     PyObject* res;
 
@@ -180,6 +182,64 @@ static PyObject* call_method(PyObject* o, const char* name, PyObject** nameobj,
     return retval;
 }
 
+
+PyObject* slot_tp_new(PyTypeObject* self, PyObject* args, PyObject* kwds) noexcept {
+    try {
+        // TODO: runtime ICs?
+        Box* new_attr = typeLookup(self, _new_str, NULL);
+        assert(new_attr);
+        new_attr = processDescriptor(new_attr, None, self);
+
+        return runtimeCall(new_attr, ArgPassSpec(1, 0, true, true), self, args, kwds, NULL, NULL);
+    } catch (Box* e) {
+        abort();
+    }
+}
+
+PyObject* slot_tp_call(PyObject* self, PyObject* args, PyObject* kwds) noexcept {
+    try {
+        Py_FatalError("this function is untested");
+
+        // TODO: runtime ICs?
+        return runtimeCall(self, ArgPassSpec(0, 0, true, true), args, kwds, NULL, NULL, NULL);
+    } catch (Box* e) {
+        abort();
+    }
+}
+
+PyObject* slot_tp_repr(PyObject* self) noexcept {
+    try {
+        return repr(self);
+    } catch (Box* e) {
+        abort();
+    }
+}
+
+PyObject* slot_sq_item(PyObject* self, Py_ssize_t i) noexcept {
+    try {
+        return getitem(self, boxInt(i));
+    } catch (Box* e) {
+        abort();
+    }
+}
+
+static Py_ssize_t slot_sq_length(PyObject* self) {
+    static PyObject* len_str;
+    PyObject* res = call_method(self, "__len__", &len_str, "()");
+    Py_ssize_t len;
+
+    if (res == NULL)
+        return -1;
+    len = PyInt_AsSsize_t(res);
+    Py_DECREF(res);
+    if (len < 0) {
+        if (!PyErr_Occurred())
+            PyErr_SetString(PyExc_ValueError, "__len__() should return >= 0");
+        return -1;
+    }
+    return len;
+}
+
 // Copied from CPython:
 #define SLOT0(FUNCNAME, OPSTR)                                                                                         \
     static PyObject* FUNCNAME(PyObject* self) noexcept {                                                               \
@@ -193,8 +253,25 @@ static PyObject* call_method(PyObject* o, const char* name, PyObject** nameobj,
         return call_method(self, OPSTR, &cache_str, "(" ARGCODES ")", arg1);                                           \
     }
 
+#define slot_mp_length slot_sq_length
+
 SLOT1(slot_mp_subscript, "__getitem__", PyObject*, "O")
 
+static int slot_mp_ass_subscript(PyObject* self, PyObject* key, PyObject* value) {
+    PyObject* res;
+    static PyObject* delitem_str, *setitem_str;
+
+    if (value == NULL)
+        res = call_method(self, "__delitem__", &delitem_str, "(O)", key);
+    else
+        res = call_method(self, "__setitem__", &setitem_str, "(OO)", key, value);
+    if (res == NULL)
+        return -1;
+    Py_DECREF(res);
+    return 0;
+}
+
+
 typedef wrapper_def slotdef;
 
 static void** slotptr(BoxedClass* type, int offset) {
@@ -266,7 +343,11 @@ static slotdef slotdefs[] = {
            PyWrapperFlag_KEYWORDS),
     TPSLOT("__new__", tp_new, slot_tp_new, NULL, ""),
 
+    MPSLOT("__len__", mp_length, slot_mp_length, wrap_lenfunc, "x.__len__() <==> len(x)"),
     MPSLOT("__getitem__", mp_subscript, slot_mp_subscript, wrap_binaryfunc, "x.__getitem__(y) <==> x[y]"),
+    MPSLOT("__setitem__", mp_ass_subscript, slot_mp_ass_subscript, wrap_objobjargproc,
+           "x.__setitem__(i, y) <==> x[i]=y"),
+    MPSLOT("__delitem__", mp_ass_subscript, slot_mp_ass_subscript, wrap_delitem, "x.__delitem__(y) <==> del x[y]"),
 
     // SQSLOT("__len__", sq_length, slot_sq_length, wrap_lenfunc, "x.__len__() <==> len(x)"),
     /* Heap types defining __add__/__mul__ have sq_concat/sq_repeat == NULL.
@@ -400,7 +481,6 @@ extern "C" int PyType_Ready(PyTypeObject* cls) {
     RELEASE_ASSERT(cls->tp_setattr == NULL, "");
     RELEASE_ASSERT(cls->tp_compare == NULL, "");
     RELEASE_ASSERT(cls->tp_as_number == NULL, "");
-    RELEASE_ASSERT(cls->tp_as_mapping == NULL, "");
     RELEASE_ASSERT(cls->tp_hash == NULL, "");
     RELEASE_ASSERT(cls->tp_str == NULL, "");
     RELEASE_ASSERT(cls->tp_getattro == NULL || cls->tp_getattro == PyObject_GenericGetAttr, "");

test/test_extension/slots_test.c

@@ -29,7 +29,7 @@ slots_tester_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
 }
 
 static PyObject *
-slots_tester_repr(slots_tester_object *obj)
+slots_tester_seq_repr(slots_tester_object *obj)
 {
     char buf[80];
     snprintf(buf, sizeof(buf), "<my custom repr: %d>", obj->n);
@@ -37,7 +37,7 @@ slots_tester_repr(slots_tester_object *obj)
 }
 
 static PyObject *
-slots_tester_call(slots_tester_object *obj, PyObject *args, PyObject *kw)
+slots_tester_seq_call(slots_tester_object *obj, PyObject *args, PyObject *kw)
 {
     if (!PyArg_ParseTuple(args, ""))
         return NULL;
@@ -46,7 +46,7 @@ slots_tester_call(slots_tester_object *obj, PyObject *args, PyObject *kw)
 }
 
 static PyObject*
-slots_tester_item(slots_tester_object *obj, Py_ssize_t i)
+slots_tester_seq_item(slots_tester_object *obj, Py_ssize_t i)
 {
     if (i < 0 || i >= 5) {
         PyErr_SetString(PyExc_IndexError, "tuple index out of range");
@@ -55,13 +55,13 @@ slots_tester_item(slots_tester_object *obj, Py_ssize_t i)
     return PyInt_FromLong(i + obj->n);
 }
 
-PyDoc_STRVAR(slots_tester_doc, "slots_tester doc");
+PyDoc_STRVAR(slots_tester_seq_doc, "slots_tester_seq doc");
 
-static PySequenceMethods slots_tester_as_sequence = {
+static PySequenceMethods slots_tester_seq_as_sequence = {
     (lenfunc)0,
     (binaryfunc)0,           /* sq_concat */
     (ssizeargfunc)0,         /* sq_repeat */
-    (ssizeargfunc)slots_tester_item,               /* sq_item */
+    (ssizeargfunc)slots_tester_seq_item,               /* sq_item */
     (ssizessizeargfunc)0,         /* sq_slice */
     0,                                          /* sq_ass_item */
     0,                                          /* sq_ass_slice */
@@ -69,9 +69,9 @@ static PySequenceMethods slots_tester_as_sequence = {
 };
 
 
-static PyTypeObject slots_tester = {
+static PyTypeObject slots_tester_seq = {
     PyVarObject_HEAD_INIT(NULL, 0)
-    "slots_test.slots_tester",            /* tp_name */
+    "slots_test.slots_tester_seq",            /* tp_name */
     sizeof(slots_tester_object),          /* tp_basicsize */
     0,                                  /* tp_itemsize */
     /* methods */
@@ -80,18 +80,83 @@ static PyTypeObject slots_tester = {
     0,                                  /* tp_getattr */
     0,                                  /* tp_setattr */
     0,                                  /* tp_compare */
-    (reprfunc)slots_tester_repr,        /* tp_repr */
+    (reprfunc)slots_tester_seq_repr,        /* tp_repr */
     0,                                  /* tp_as_number */
-    &slots_tester_as_sequence,          /* tp_as_sequence */
+    &slots_tester_seq_as_sequence,          /* tp_as_sequence */
     0,                                  /* tp_as_mapping */
     0,                                  /* tp_hash */
-    (ternaryfunc)slots_tester_call,     /* tp_call */
+    (ternaryfunc)slots_tester_seq_call,     /* tp_call */
     0,                                  /* tp_str */
     0,                                  /* tp_getattro */
     0,                                  /* tp_setattro */
     0,                                  /* tp_as_buffer */
     Py_TPFLAGS_DEFAULT,                 /* tp_flags */
-    slots_tester_doc,                   /* tp_doc */
+    slots_tester_seq_doc,                   /* tp_doc */
+    0,                                  /* tp_traverse */
+    0,                                  /* tp_clear */
+    0,                                  /* tp_richcompare */
+    0,                                  /* tp_weaklistoffset */
+    0,                                  /* tp_iter */
+    0,                                  /* 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 */
+    0,                                  /* tp_init */
+    0,                                  /* tp_alloc */
+    slots_tester_new,                   /* tp_new */
+    0,                                  /* tp_free */
+};
+
+static Py_ssize_t slots_tester_map_length(slots_tester_object* obj) {
+    return obj->n;
+}
+
+static PyObject* slots_tester_map_subscript(slots_tester_object* obj, PyObject* key) {
+    int n2 = PyInt_AsLong(key);
+    return PyInt_FromLong(n2 + obj->n);
+}
+
+static int slots_tester_map_ass_sub(slots_tester_object* obj, PyObject* key, PyObject* value) {
+    int n2 = PyInt_AsLong(key);
+    printf("Assigning to subscript for object with n=%d, key=%d, set/delete=%d\n", obj->n, n2, !!value);
+
+    return 0;
+}
+
+static PyMappingMethods slots_tester_map_asmapping = {
+    (lenfunc)slots_tester_map_length,              /*mp_length*/
+    (binaryfunc)slots_tester_map_subscript,        /*mp_subscript*/
+    (objobjargproc)slots_tester_map_ass_sub,       /*mp_ass_subscript*/
+};
+
+static PyTypeObject slots_tester_map= {
+    PyVarObject_HEAD_INIT(NULL, 0)
+    "slots_test.slots_tester_map",            /* tp_name */
+    sizeof(slots_tester_object),          /* tp_basicsize */
+    0,                                  /* tp_itemsize */
+    /* methods */
+    0,                                  /* tp_dealloc */
+    0,                                  /* tp_print */
+    0,                                  /* tp_getattr */
+    0,                                  /* tp_setattr */
+    0,                                  /* tp_compare */
+    0,        /* tp_repr */
+    0,                                  /* tp_as_number */
+    0,          /* tp_as_sequence */
+    &slots_tester_map_asmapping,                                  /* 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,                 /* tp_flags */
+    0,                   /* tp_doc */
     0,                                  /* tp_traverse */
     0,                                  /* tp_clear */
     0,                                  /* tp_richcompare */
@@ -185,12 +250,17 @@ initslots_test(void)
     if (m == NULL)
         return;
 
-    int res = PyType_Ready(&slots_tester);
+    int res = PyType_Ready(&slots_tester_seq);
+    if (res < 0)
+        return;
+
+    res = PyType_Ready(&slots_tester_map);
     if (res < 0)
         return;
 
     // Not sure if the result of PyInt_FromLong needs to be decref'd
-    PyDict_SetItemString(slots_tester.tp_dict, "set_through_tpdict", PyInt_FromLong(123));
+    PyDict_SetItemString(slots_tester_seq.tp_dict, "set_through_tpdict", PyInt_FromLong(123));
 
-    PyModule_AddObject(m, "SlotsTester", (PyObject *)&slots_tester);
+    PyModule_AddObject(m, "SlotsTesterSeq", (PyObject *)&slots_tester_seq);
+    PyModule_AddObject(m, "SlotsTesterMap", (PyObject *)&slots_tester_map);
 }

test/tests/capi_slots.py

@@ -4,10 +4,17 @@
 import slots_test
 
 for i in xrange(3):
-    t = slots_test.SlotsTester(i + 5)
+    t = slots_test.SlotsTesterSeq(i + 5)
     print t, repr(t), t(), t[2]
 
-print slots_test.SlotsTester.set_through_tpdict, slots_test.SlotsTester(5).set_through_tpdict
+# print slots_test.SlotsTesterSeq.__doc__
+print slots_test.SlotsTesterSeq.set_through_tpdict, slots_test.SlotsTesterSeq(5).set_through_tpdict
+
+for i in xrange(3):
+    t = slots_test.SlotsTesterMap(i + 5)
+    print len(t), t[2]
+    t[1] = 5
+    del t[2]
 
 class C(object):
     def __repr__(self):