optimise dict.setdefault() - would be better in Builtin.py, but class scopes...

optimise dict.setdefault() - would be better in Builtin.py, but class scopes don't seem to support signature overloading yet

--HG--
extra : rebase_source : d0f2bd97e0cec8919ac00e9a0ec8b714e338a9cb

Cython/Compiler/Optimize.py

@@ -2359,6 +2359,28 @@ class OptimizeBuiltinCalls(Visitor.EnvTransform):
             may_return_none = True,
             utility_code = dict_getitem_default_utility_code)
 
+    Pyx_PyDict_SetDefault_func_type = PyrexTypes.CFuncType(
+        PyrexTypes.py_object_type, [
+            PyrexTypes.CFuncTypeArg("dict", PyrexTypes.py_object_type, None),
+            PyrexTypes.CFuncTypeArg("key", PyrexTypes.py_object_type, None),
+            PyrexTypes.CFuncTypeArg("default", PyrexTypes.py_object_type, None),
+            ])
+
+    def _handle_simple_method_dict_setdefault(self, node, args, is_unbound_method):
+        """Replace dict.setdefault() by calls to PyDict_GetItem() and PyDict_SetItem().
+        """
+        if len(args) == 2:
+            args.append(ExprNodes.NoneNode(node.pos))
+        elif len(args) != 3:
+            self._error_wrong_arg_count('dict.setdefault', node, args, "2 or 3")
+            return node
+
+        return self._substitute_method_call(
+            node, "__Pyx_PyDict_SetDefault", self.Pyx_PyDict_SetDefault_func_type,
+            'setdefault', is_unbound_method, args,
+            may_return_none = True,
+            utility_code = dict_setdefault_utility_code)
+
 
     ### unicode type methods
 
@@ -3103,6 +3125,45 @@ static PyObject* __Pyx_PyDict_GetItemDefault(PyObject* d, PyObject* key, PyObjec
 impl = ""
 )
 
+dict_setdefault_utility_code = UtilityCode(
+proto = """
+static PyObject *__Pyx_PyDict_SetDefault(PyObject *, PyObject *, PyObject *); /*proto*/
+""",
+impl = '''
+static PyObject *__Pyx_PyDict_SetDefault(PyObject *d, PyObject *key, PyObject *default_value) {
+    PyObject* value;
+#if PY_MAJOR_VERSION >= 3
+    value = PyDict_GetItemWithError(d, key);
+    if (unlikely(!value)) {
+        if (unlikely(PyErr_Occurred()))
+            return NULL;
+        if (unlikely(PyDict_SetItem(d, key, default_value) == -1))
+            return NULL;
+        value = default_value;
+    }
+    Py_INCREF(value);
+#else
+    if (PyString_CheckExact(key) || PyUnicode_CheckExact(key) || PyInt_CheckExact(key)) {
+        /* these presumably have safe hash functions */
+        value = PyDict_GetItem(d, key);
+        if (unlikely(!value)) {
+            if (unlikely(PyDict_SetItem(d, key, default_value) == -1))
+                return NULL;
+            value = default_value;
+        }
+        Py_INCREF(value);
+    } else {
+        PyObject *m;
+        m = __Pyx_GetAttrString(d, "setdefault");
+        if (!m) return NULL;
+        value = PyObject_CallFunctionObjArgs(m, key, default_value, NULL);
+        Py_DECREF(m);
+    }
+#endif
+    return value;
+}
+''')
+
 append_utility_code = UtilityCode(
 proto = """
 static CYTHON_INLINE PyObject* __Pyx_PyObject_Append(PyObject* L, PyObject* x) {

tests/run/dict_setdefault.py

+
+import cython
+
+class Unhashable(object):
+    def __hash__(self):
+        raise TypeError('I am not hashable')
+
+class Hashable(object):
+    def __hash__(self):
+        return 1
+    def __eq__(self, other):
+        return isinstance(other, Hashable)
+
+@cython.test_fail_if_path_exists('//AttributeNode')
+@cython.test_assert_path_exists('//PythonCapiCallNode')
+@cython.locals(d=dict)
+def setdefault1(d, key):
+    """
+    >>> d = {}
+    >>> setdefault1(d, 1)
+    >>> len(d)
+    1
+    >>> setdefault1(d, 1)
+    >>> len(d)
+    1
+    >>> d[1]
+    >>> setdefault1(d, Unhashable())
+    Traceback (most recent call last):
+    TypeError: I am not hashable
+    >>> len(d)
+    1
+    >>> h1 = setdefault1(d, Hashable())
+    >>> len(d)
+    2
+    >>> h2 = setdefault1(d, Hashable())
+    >>> len(d)
+    2
+    >>> d[Hashable()]
+    """
+    return d.setdefault(key)
+
+@cython.test_fail_if_path_exists('//AttributeNode')
+@cython.test_assert_path_exists('//PythonCapiCallNode')
+@cython.locals(d=dict)
+def setdefault2(d, key, value):
+    """
+    >>> d = {}
+    >>> setdefault2(d, 1, 2)
+    2
+    >>> len(d)
+    1
+    >>> setdefault2(d, 1, 2)
+    2
+    >>> len(d)
+    1
+    >>> l = setdefault2(d, 2, [])
+    >>> len(d)
+    2
+    >>> l.append(1)
+    >>> setdefault2(d, 2, [])
+    [1]
+    >>> len(d)
+    2
+    >>> setdefault2(d, Unhashable(), 1)
+    Traceback (most recent call last):
+    TypeError: I am not hashable
+    >>> h1 = setdefault2(d, Hashable(), 55)
+    >>> len(d)
+    3
+    >>> h2 = setdefault2(d, Hashable(), 66)
+    >>> len(d)
+    3
+    >>> d[Hashable()]
+    55
+    """
+    return d.setdefault(key, value)