Merge pull request #2788 from cython/gh1807_getitem_mapping_first

Try Mapping protocol before Sequence protocol on getitem/setitem/delitem

Cython/Utility/ObjectHandling.c

@@ -429,10 +429,15 @@ static CYTHON_INLINE PyObject *__Pyx_GetItemInt_Fast(PyObject *o, Py_ssize_t i,
         }
     } else {
         // inlined PySequence_GetItem() + special cased length overflow
-        PySequenceMethods *m = Py_TYPE(o)->tp_as_sequence;
-        if (likely(m && m->sq_item)) {
-            if (wraparound && unlikely(i < 0) && likely(m->sq_length)) {
-                Py_ssize_t l = m->sq_length(o);
+        PyMappingMethods *mm = Py_TYPE(o)->tp_as_mapping;
+        PySequenceMethods *sm = Py_TYPE(o)->tp_as_sequence;
+        if (mm && mm->mp_subscript) {
+            PyObject *key = PyInt_FromSsize_t(i);
+            return likely(key) ? mm->mp_subscript(o, key) : NULL;
+        }
+        if (likely(sm && sm->sq_item)) {
+            if (wraparound && unlikely(i < 0) && likely(sm->sq_length)) {
+                Py_ssize_t l = sm->sq_length(o);
                 if (likely(l >= 0)) {
                     i += l;
                 } else {
@@ -442,7 +447,7 @@ static CYTHON_INLINE PyObject *__Pyx_GetItemInt_Fast(PyObject *o, Py_ssize_t i,
                     PyErr_Clear();
                 }
             }
-            return m->sq_item(o, i);
+            return sm->sq_item(o, i);
         }
     }
 #else
@@ -489,10 +494,15 @@ static CYTHON_INLINE int __Pyx_SetItemInt_Fast(PyObject *o, Py_ssize_t i, PyObje
         }
     } else {
         // inlined PySequence_SetItem() + special cased length overflow
-        PySequenceMethods *m = Py_TYPE(o)->tp_as_sequence;
-        if (likely(m && m->sq_ass_item)) {
-            if (wraparound && unlikely(i < 0) && likely(m->sq_length)) {
-                Py_ssize_t l = m->sq_length(o);
+        PyMappingMethods *mm = Py_TYPE(o)->tp_as_mapping;
+        PySequenceMethods *sm = Py_TYPE(o)->tp_as_sequence;
+        if (mm && mm->mp_ass_subscript) {
+            PyObject *key = PyInt_FromSsize_t(i);
+            return likely(key) ? mm->mp_ass_subscript(o, key, v) : -1;
+        }
+        if (likely(sm && sm->sq_ass_item)) {
+            if (wraparound && unlikely(i < 0) && likely(sm->sq_length)) {
+                Py_ssize_t l = sm->sq_length(o);
                 if (likely(l >= 0)) {
                     i += l;
                 } else {
@@ -502,7 +512,7 @@ static CYTHON_INLINE int __Pyx_SetItemInt_Fast(PyObject *o, Py_ssize_t i, PyObje
                     PyErr_Clear();
                 }
             }
-            return m->sq_ass_item(o, i, v);
+            return sm->sq_ass_item(o, i, v);
         }
     }
 #else
@@ -542,17 +552,22 @@ static int __Pyx_DelItem_Generic(PyObject *o, PyObject *j) {
 }
 
 static CYTHON_INLINE int __Pyx_DelItemInt_Fast(PyObject *o, Py_ssize_t i,
-                                               CYTHON_UNUSED int is_list, CYTHON_NCP_UNUSED int wraparound) {
+                                               int is_list, CYTHON_NCP_UNUSED int wraparound) {
 #if !CYTHON_USE_TYPE_SLOTS
     if (is_list || PySequence_Check(o)) {
         return PySequence_DelItem(o, i);
     }
 #else
     // inlined PySequence_DelItem() + special cased length overflow
-    PySequenceMethods *m = Py_TYPE(o)->tp_as_sequence;
-    if (likely(m && m->sq_ass_item)) {
-        if (wraparound && unlikely(i < 0) && likely(m->sq_length)) {
-            Py_ssize_t l = m->sq_length(o);
+    PyMappingMethods *mm = Py_TYPE(o)->tp_as_mapping;
+    PySequenceMethods *sm = Py_TYPE(o)->tp_as_sequence;
+    if ((!is_list) && mm && mm->mp_ass_subscript) {
+        PyObject *key = PyInt_FromSsize_t(i);
+        return likely(key) ? mm->mp_ass_subscript(o, key, (PyObject *)NULL) : -1;
+    }
+    if (likely(sm && sm->sq_ass_item)) {
+        if (wraparound && unlikely(i < 0) && likely(sm->sq_length)) {
+            Py_ssize_t l = sm->sq_length(o);
             if (likely(l >= 0)) {
                 i += l;
             } else {
@@ -562,7 +577,7 @@ static CYTHON_INLINE int __Pyx_DelItemInt_Fast(PyObject *o, Py_ssize_t i,
                 PyErr_Clear();
             }
         }
-        return m->sq_ass_item(o, i, (PyObject *)NULL);
+        return sm->sq_ass_item(o, i, (PyObject *)NULL);
     }
 #endif
     return __Pyx_DelItem_Generic(o, PyInt_FromSsize_t(i));

tests/run/dict_getitem.pyx

@@ -146,3 +146,17 @@ def getitem_not_none(dict d not None, key):
     KeyError: (1, 2)
     """
     return d[key]
+
+
+def getitem_int_key(d, int key):
+    """
+    >>> d = {-1: 10}
+    >>> getitem_int_key(d, -1)  # dict
+    10
+    >>> class D(dict): pass
+    >>> d = D({-1: 10})
+    >>> getitem_int_key(d, -1)  # D
+    10
+    """
+    # Based on GH-1807: must check Mapping protocol first, even for integer "index" keys.
+    return d[key]

tests/run/set_item.pyx

+# mode: run
+# tag: list, dict, setitem, delitem
+
+def set_item(obj, key, value):
+    """
+    >>> set_item([1, 2, 3], 1, -1)
+    [1, -1, 3]
+    >>> set_item([1, 2, 3], -1, -1)
+    [1, 2, -1]
+    >>> set_item({}, 'abc', 5)
+    {'abc': 5}
+    >>> set_item({}, -1, 5)
+    {-1: 5}
+    >>> class D(dict): pass
+    >>> set_item(D({}), 'abc', 5)
+    {'abc': 5}
+    >>> set_item(D({}), -1, 5)
+    {-1: 5}
+    """
+    obj[key] = value
+    return obj
+
+
+def set_item_int(obj, int key, value):
+    """
+    >>> set_item_int([1, 2, 3], 1, -1)
+    [1, -1, 3]
+    >>> set_item_int([1, 2, 3], -1, -1)
+    [1, 2, -1]
+    >>> set_item_int({}, 1, 5)
+    {1: 5}
+    >>> set_item_int({}, -1, 5)
+    {-1: 5}
+    >>> class D(dict): pass
+    >>> set_item_int(D({}), 1, 5)
+    {1: 5}
+    >>> set_item_int(D({}), -1, 5)
+    {-1: 5}
+    """
+    obj[key] = value
+    return obj
+
+
+def del_item(obj, key):
+    """
+    >>> del_item([1, 2, 3], 1)
+    [1, 3]
+    >>> del_item([1, 2, 3], -3)
+    [2, 3]
+    >>> class D(dict): pass
+    >>> del_item({'abc': 1, 'def': 2}, 'abc')
+    {'def': 2}
+    >>> del_item(D({'abc': 1, 'def': 2}), 'abc')
+    {'def': 2}
+    >>> del_item(D({-1: 1, -2: 2}), -1)
+    {-2: 2}
+    """
+    del obj[key]
+    return obj
+
+
+def del_item_int(obj, int key):
+    """
+    >>> del_item_int([1, 2, 3], 1)
+    [1, 3]
+    >>> del_item_int([1, 2, 3], -3)
+    [2, 3]
+    >>> class D(dict): pass
+    >>> del_item_int(D({-1: 1, 1: 2}), 1)
+    {-1: 1}
+    >>> del_item_int(D({-1: 1, -2: 2}), -1)
+    {-2: 2}
+    """
+    del obj[key]
+    return obj