diff --git a/mypyc/lib-rt/dict_ops.c b/mypyc/lib-rt/dict_ops.c index 8a9202c403f8..8894ae18b92c 100644 --- a/mypyc/lib-rt/dict_ops.c +++ b/mypyc/lib-rt/dict_ops.c @@ -212,58 +212,35 @@ PyObject *CPyDict_ItemsView(PyObject *dict) { return PyObject_CallMethodNoArgs(dict, mypyc_interned_str.items); } -PyObject *CPyDict_Keys(PyObject *dict) { - if (PyDict_CheckExact(dict)) { - return PyDict_Keys(dict); - } - // Inline generic fallback logic to also return a list. - PyObject *list = PyList_New(0); - PyObject *view = PyObject_CallMethodNoArgs(dict, mypyc_interned_str.keys); +static inline PyObject *dict_view_as_list(PyObject *dict, PyObject *method) { + PyObject *view = PyObject_CallMethodNoArgs(dict, method); if (view == NULL) { return NULL; } - int res = PyList_Extend(list, view); + PyObject *list = PySequence_List(view); Py_DECREF(view); - if (res < 0) { - return NULL; - } return list; } +PyObject *CPyDict_Keys(PyObject *dict) { + if (PyDict_CheckExact(dict)) { + return PyDict_Keys(dict); + } + return dict_view_as_list(dict, mypyc_interned_str.keys); +} + PyObject *CPyDict_Values(PyObject *dict) { if (PyDict_CheckExact(dict)) { return PyDict_Values(dict); } - // Inline generic fallback logic to also return a list. - PyObject *list = PyList_New(0); - PyObject *view = PyObject_CallMethodNoArgs(dict, mypyc_interned_str.values); - if (view == NULL) { - return NULL; - } - int res = PyList_Extend(list, view); - Py_DECREF(view); - if (res < 0) { - return NULL; - } - return list; + return dict_view_as_list(dict, mypyc_interned_str.values); } PyObject *CPyDict_Items(PyObject *dict) { if (PyDict_CheckExact(dict)) { return PyDict_Items(dict); } - // Inline generic fallback logic to also return a list. - PyObject *list = PyList_New(0); - PyObject *view = PyObject_CallMethodNoArgs(dict, mypyc_interned_str.items); - if (view == NULL) { - return NULL; - } - int res = PyList_Extend(list, view); - Py_DECREF(view); - if (res < 0) { - return NULL; - } - return list; + return dict_view_as_list(dict, mypyc_interned_str.items); } char CPyDict_Clear(PyObject *dict) { diff --git a/mypyc/test-data/run-dicts.test b/mypyc/test-data/run-dicts.test index da507900ff50..9cf57fbfc555 100644 --- a/mypyc/test-data/run-dicts.test +++ b/mypyc/test-data/run-dicts.test @@ -513,3 +513,69 @@ test_sub_cyclic_dict_cleanup() test_cyclic_dict_cleanup_interpreted() test_sub_cyclic_dict_cleanup_interpreted() test_allow_interpreted_subclass_cyclic_dict_cleanup() + +[case testDictListErrorOwnership] +from typing import Any, Dict, List + +def keys(d: Dict[Any, Any]) -> List[Any]: + return list(d.keys()) +def values(d: Dict[Any, Any]) -> List[Any]: + return list(d.values()) +def items(d: Dict[Any, Any]) -> List[Any]: + return list(d.items()) + +[file driver.py] +import gc +import weakref +from native import keys, values, items +from testutil import assertRaises + +class Sentinel: + pass + +refs = [] +def failing_iterator(): + value = Sentinel() + refs.append(weakref.ref(value)) + yield value + raise RuntimeError("iteration failed") + +class BadDict(dict): + def keys(self): + return failing_iterator() + values = keys + items = keys + +class BadMethod(dict): + def keys(self): + raise RuntimeError("method failed") + values = keys + items = keys + +class BadIter: + def __iter__(self): + raise RuntimeError("iter failed") + +class BadLength: + def __iter__(self): + return self + def __next__(self): + return None + def __length_hint__(self): + raise RuntimeError("length failed") + +class CustomDict(dict): + def __init__(self, result): + self.result = result + def keys(self): + return self.result + values = keys + items = keys + +for convert in (keys, values, items): + for d in (BadDict(), BadMethod(), CustomDict(BadIter()), CustomDict(BadLength())): + with assertRaises(RuntimeError): + convert(d) + gc.collect() + assert all(ref() is None for ref in refs) + assert convert(CustomDict(iter([1, 2]))) == [1, 2]