Tune dict lookups with Unicode string keys.

Cython/Compiler/ExprNodes.py

@@ -3883,6 +3883,8 @@ class IndexNode(_IndexingBaseNode):
         if not self.is_temp:
             # all handled in self.calculate_result_code()
             return
+
+        utility_code = None
         if self.type.is_pyobject:
             error_value = 'NULL'
             if self.index.type.is_int:
@@ -3892,32 +3894,39 @@ class IndexNode(_IndexingBaseNode):
                     function = "__Pyx_GetItemInt_Tuple"
                 else:
                     function = "__Pyx_GetItemInt"
-                code.globalstate.use_utility_code(
-                    TempitaUtilityCode.load_cached("GetItemInt", "ObjectHandling.c"))
+                utility_code = TempitaUtilityCode.load_cached("GetItemInt", "ObjectHandling.c")
             else:
                 if self.base.type is dict_type:
-                    function = "__Pyx_PyDict_GetItem"
-                    code.globalstate.use_utility_code(
-                        UtilityCode.load_cached("DictGetItem", "ObjectHandling.c"))
+                    if self.index.type in (str_type, unicode_type):
+                        function = "__Pyx_PyDict_GetItemUnicode"
+                        utility_code = UtilityCode.load_cached("DictGetItemUnicode", "ObjectHandling.c")
+                    else:
+                        function = "__Pyx_PyDict_GetItem"
+                        utility_code = UtilityCode.load_cached("DictGetItem", "ObjectHandling.c")
+                elif self.index.type in (str_type, unicode_type):
+                    # obj[str] is probably doing a dict lookup
+                    function = "__Pyx_PyObject_GetItemUnicode"
+                    utility_code = UtilityCode.load_cached("DictGetItemUnicode", "ObjectHandling.c")
                 else:
                     function = "PyObject_GetItem"
         elif self.type.is_unicode_char and self.base.type is unicode_type:
             assert self.index.type.is_int
             function = "__Pyx_GetItemInt_Unicode"
             error_value = '(Py_UCS4)-1'
-            code.globalstate.use_utility_code(
-                UtilityCode.load_cached("GetItemIntUnicode", "StringTools.c"))
+            utility_code = UtilityCode.load_cached("GetItemIntUnicode", "StringTools.c")
         elif self.base.type is bytearray_type:
             assert self.index.type.is_int
             assert self.type.is_int
             function = "__Pyx_GetItemInt_ByteArray"
             error_value = '-1'
-            code.globalstate.use_utility_code(
-                UtilityCode.load_cached("GetItemIntByteArray", "StringTools.c"))
+            utility_code = UtilityCode.load_cached("GetItemIntByteArray", "StringTools.c")
         elif not (self.base.type.is_cpp_class and self.exception_check):
             assert False, "unexpected type %s and base type %s for indexing" % (
                 self.type, self.base.type)
 
+        if utility_code is not None:
+            code.globalstate.use_utility_code(utility_code)
+
         if self.index.type.is_int:
             index_code = self.index.result()
         else:

Cython/Utility/ObjectHandling.c

@@ -262,8 +262,59 @@ static CYTHON_INLINE int __Pyx_IterFinish(void) {
 #endif
 }
 
+/////////////// DictGetItemUnicode.proto ///////////////
+//@requires: DictGetItem
+
+#if CYTHON_COMPILING_IN_CPYTHON && PY_VERSION_HEX >= 0x030500A1 && CYTHON_USE_UNICODE_INTERNALS
+static CYTHON_INLINE PyObject *__Pyx__PyDict_GetItemUnicode(PyObject *d, PyObject* key);
+
+#define __Pyx_PyDict_GetItemUnicode(d, key)  \
+    (likely(PyUnicode_CheckExact(key) && ((PyASCIIObject *) key)->hash != -1) ? \
+     __Pyx__PyDict_GetItemUnicode(d, key) : __Pyx_PyDict_GetItem(d, key))
+
+#define __Pyx_PyObject_GetItemUnicode(obj, name) \
+    (likely(PyDict_CheckExact(obj)) ? \
+     __Pyx_PyDict_GetItemUnicode(obj, name) : PyObject_GetItem(obj, name))
+
+#else /* Py < 3.5 */
+#define __Pyx_PyDict_GetItemUnicode(d, key)  __Pyx_PyDict_GetItem(d, key)
+
+#if PY_MAJOR_VERSION >= 3 && !CYTHON_COMPILING_IN_PYPY
+#define __Pyx_PyObject_GetItemUnicode(obj, name) \
+    (likely(PyDict_CheckExact(obj)) ? \
+     __Pyx_PyDict_GetItem(obj, name) : PyObject_GetItem(obj, name))
+#else
+#define __Pyx_PyObject_GetItemUnicode(obj, name)  PyObject_GetItem(obj, name)
+#endif
+
+#endif
+
+/////////////// DictGetItemUnicode ///////////////
+
+#if CYTHON_COMPILING_IN_CPYTHON && PY_VERSION_HEX >= 0x030500A1 && CYTHON_USE_UNICODE_INTERNALS
+static PyObject *__Pyx__PyDict_GetItemUnicode(PyObject *d, PyObject* key) {
+    PyObject *value;
+    value = __Pyx_PyDict_GetItemStr(d, key);
+    if (unlikely(!value)) {
+        PyErr_SetObject(PyExc_KeyError, key);
+        return NULL;
+    }
+    Py_INCREF(value);
+    return value;
+}
+#endif
+
+
 /////////////// DictGetItem.proto ///////////////
 
+#if PY_MAJOR_VERSION >= 3 && !CYTHON_COMPILING_IN_PYPY
+static PyObject *__Pyx_PyDict_GetItem(PyObject *d, PyObject* key);/*proto*/
+#else
+#define __Pyx_PyDict_GetItem(d, key) PyObject_GetItem(d, key)
+#endif
+
+/////////////// DictGetItem ///////////////
+
 #if PY_MAJOR_VERSION >= 3 && !CYTHON_COMPILING_IN_PYPY
 static PyObject *__Pyx_PyDict_GetItem(PyObject *d, PyObject* key) {
     PyObject *value;
@@ -280,8 +331,6 @@ static PyObject *__Pyx_PyDict_GetItem(PyObject *d, PyObject* key) {
     Py_INCREF(value);
     return value;
 }
-#else
-    #define __Pyx_PyDict_GetItem(d, key) PyObject_GetItem(d, key)
 #endif
 
 /////////////// GetItemInt.proto ///////////////

tests/run/dict_getitem.pyx

@@ -46,6 +46,72 @@ def test(dict d, index):
     """
     return d[index]
 
+
+def getitem_str(dict d, obj, str index):
+    """
+    >>> d = {'abc': 1, 'xyz': 2, None: 3}
+    >>> getitem_str(d, d, 'abc')
+    (1, 1)
+    >>> getitem_str(d, d, 'xyz')
+    (2, 2)
+    >>> getitem_str(d, d, None)
+    (3, 3)
+
+    >>> class GetItem(object):
+    ...     def __getitem__(self, name): return d[name]
+    >>> getitem_str(d, GetItem(), 'abc')
+    (1, 1)
+    >>> getitem_str(d, GetItem(), 'xyz')
+    (2, 2)
+    >>> getitem_str(d, GetItem(), None)
+    (3, 3)
+    >>> getitem_str(d, GetItem(), 'no')
+    Traceback (most recent call last):
+    KeyError: 'no'
+
+    >>> class GetItemFail(object):
+    ...     def __getitem__(self, name): raise ValueError("failed")
+    >>> getitem_str(d, GetItemFail(), 'abc')
+    Traceback (most recent call last):
+    ValueError: failed
+    >>> getitem_str(d, GetItemFail(), None)
+    Traceback (most recent call last):
+    ValueError: failed
+    """
+    return d[index], obj[index]
+
+
+def getitem_unicode(dict d, obj, unicode index):
+    """
+    >>> d = {'abc': 1, 'xyz': 2, None: 3}
+    >>> getitem_unicode(d, d, u'abc')
+    (1, 1)
+    >>> getitem_unicode(d, d, u'xyz')
+    (2, 2)
+    >>> getitem_unicode(d, d, None)
+    (3, 3)
+
+    >>> class GetItem(object):
+    ...     def __getitem__(self, name): return d[name]
+    >>> getitem_unicode(d, GetItem(), u'abc')
+    (1, 1)
+    >>> getitem_unicode(d, GetItem(), u'xyz')
+    (2, 2)
+    >>> getitem_unicode(d, GetItem(), None)
+    (3, 3)
+    >>> try: getitem_unicode(d, GetItem(), u'no')
+    ... except KeyError as exc: assert exc.args[0] == u'no', str(exc)
+    ... else: assert False, "KeyError not raised"
+
+    >>> class GetItemFail(object):
+    ...     def __getitem__(self, name): raise ValueError("failed")
+    >>> getitem_unicode(d, GetItemFail(), u'abc')
+    Traceback (most recent call last):
+    ValueError: failed
+    """
+    return d[index], obj[index]
+
+
 def getitem_tuple(dict d, index):
     """
     >>> d = {1: 1, (1,): 2}
@@ -54,6 +120,7 @@ def getitem_tuple(dict d, index):
     """
     return d[index], d[index,]
 
+
 def getitem_in_condition(dict d, key, expected_result):
     """
     >>> d = dict(a=1, b=2)
@@ -62,6 +129,7 @@ def getitem_in_condition(dict d, key, expected_result):
     """
     return d[key] is expected_result or d[key] == expected_result
 
+
 @cython.test_fail_if_path_exists('//NoneCheckNode')
 def getitem_not_none(dict d not None, key):
     """