diff --git a/Lib/test/test_frame.py b/Lib/test/test_frame.py index b7ef6cefaabbc0..ca88e657367d9a 100644 --- a/Lib/test/test_frame.py +++ b/Lib/test/test_frame.py @@ -15,6 +15,7 @@ from test import support from test.support import import_helper, threading_helper from test.support.script_helper import assert_python_ok +from test import mapping_tests class ClearTest(unittest.TestCase): @@ -431,6 +432,132 @@ def test_is_mapping(self): kind = "other" self.assertEqual(kind, "mapping") + def _x_stringlikes(self): + class StringSubclass(str): + pass + + class ImpostorX: + def __hash__(self): + return hash('x') + + def __eq__(self, other): + return other == 'x' + + return StringSubclass('x'), ImpostorX(), 'x' + + def test_proxy_key_stringlikes_overwrite(self): + def f(obj): + x = 1 + proxy = sys._getframe().f_locals + proxy[obj] = 2 + return ( + list(proxy.keys()), + dict(proxy), + proxy + ) + + for obj in self._x_stringlikes(): + with self.subTest(cls=type(obj).__name__): + + keys_snapshot, proxy_snapshot, proxy = f(obj) + expected_keys = ['obj', 'x', 'proxy'] + expected_dict = {'obj': 'x', 'x': 2, 'proxy': proxy} + self.assertEqual(proxy.keys(), expected_keys) + self.assertEqual(proxy, expected_dict) + self.assertEqual(keys_snapshot, expected_keys) + self.assertEqual(proxy_snapshot, expected_dict) + + def test_proxy_key_stringlikes_ftrst_write(self): + def f(obj): + proxy = sys._getframe().f_locals + proxy[obj] = 2 + self.assertEqual(x, 2) + x = 1 + + for obj in self._x_stringlikes(): + with self.subTest(cls=type(obj).__name__): + f(obj) + + def test_proxy_key_unhashables(self): + class StringSubclass(str): + __hash__ = None + + class ObjectSubclass: + __hash__ = None + + proxy = sys._getframe().f_locals + + for obj in StringSubclass('x'), ObjectSubclass(): + with self.subTest(cls=type(obj).__name__): + with self.assertRaises(TypeError): + proxy[obj] + with self.assertRaises(TypeError): + proxy[obj] = 0 + + +class FrameLocalsProxyMappingTests(mapping_tests.TestHashMappingProtocol): + """Test that FrameLocalsProxy behaves like a Mapping (with exceptions)""" + + def _f(*args, **kwargs): + def _f(): + return sys._getframe().f_locals + return _f() + type2test = _f + + @unittest.skipIf(True, 'Locals proxies for different frames never compare as equal') + def test_constructor(self): + pass + + @unittest.skipIf(True, 'Unlike a mapping: del proxy[key] fails') + def test_write(self): + pass + + @unittest.skipIf(True, 'Unlike a mapping: no proxy.popitem') + def test_popitem(self): + pass + + @unittest.skipIf(True, 'Unlike a mapping: no proxy.pop') + def test_pop(self): + pass + + @unittest.skipIf(True, 'Unlike a mapping: no proxy.clear') + def test_clear(self): + pass + + @unittest.skipIf(True, 'Unlike a mapping: no proxy.fromkeys') + def test_fromkeys(self): + pass + + # no del + def test_getitem(self): + mapping_tests.BasicTestMappingProtocol.test_getitem(self) + d = self._full_mapping({'a': 1, 'b': 2}) + self.assertEqual(d['a'], 1) + self.assertEqual(d['b'], 2) + d['c'] = 3 + d['a'] = 4 + self.assertEqual(d['c'], 3) + self.assertEqual(d['a'], 4) + + @unittest.skipIf(True, 'Unlike a mapping: no proxy.update') + def test_update(self): + pass + + # proxy.copy returns a regular dict + def test_copy(self): + d = self._full_mapping({1:1, 2:2, 3:3}) + self.assertEqual(d.copy(), {1:1, 2:2, 3:3}) + d = self._empty_mapping() + self.assertEqual(d.copy(), d) + self.assertRaises(TypeError, d.copy, None) + + self.assertIsInstance(d.copy(), dict) + + @unittest.skipIf(True, 'Locals proxies for different frames never compare as equal') + def test_eq(self): + pass + + class TestFrameCApi(unittest.TestCase): def test_basic(self): x = 1 diff --git a/Misc/NEWS.d/next/Core_and_Builtins/2024-07-26-13-56-32.gh-issue-120906.qBh2I9.rst b/Misc/NEWS.d/next/Core_and_Builtins/2024-07-26-13-56-32.gh-issue-120906.qBh2I9.rst new file mode 100644 index 00000000000000..2b753bc37d4a39 --- /dev/null +++ b/Misc/NEWS.d/next/Core_and_Builtins/2024-07-26-13-56-32.gh-issue-120906.qBh2I9.rst @@ -0,0 +1 @@ +:attr:`frame.f_locals` now supports arbitrary hashable objects as keys. diff --git a/Objects/frameobject.c b/Objects/frameobject.c index 88093eb9071ae4..a8be7d75371c16 100644 --- a/Objects/frameobject.c +++ b/Objects/frameobject.c @@ -53,22 +53,27 @@ static int framelocalsproxy_getkeyindex(PyFrameObject *frame, PyObject* key, bool read) { /* - * Returns the fast locals index of the key + * Returns -2 (!) if an error occurred; exception will be set. + * Returns the fast locals index of the key on success: * - if read == true, returns the index if the value is not NULL * - if read == false, returns the index if the value is not hidden + * Otherwise returns -1. */ - assert(PyUnicode_CheckExact(key)); - PyCodeObject *co = _PyFrame_GetCode(frame->f_frame); - int found_key = false; + + // Ensure that the key is hashable. + Py_hash_t key_hash = PyObject_Hash(key); + if (key_hash == -1) { + return -2; + } + bool found = false; // We do 2 loops here because it's highly possible the key is interned // and we can do a pointer comparison. for (int i = 0; i < co->co_nlocalsplus; i++) { PyObject *name = PyTuple_GET_ITEM(co->co_localsplusnames, i); if (name == key) { - found_key = true; if (read) { if (framelocalsproxy_getval(frame->f_frame, co, i) != NULL) { return i; @@ -78,23 +83,35 @@ framelocalsproxy_getkeyindex(PyFrameObject *frame, PyObject* key, bool read) return i; } } + found = true; } } - - if (!found_key) { - // This is unlikely, but we need to make sure. This means the key - // is not interned. - for (int i = 0; i < co->co_nlocalsplus; i++) { - PyObject *name = PyTuple_GET_ITEM(co->co_localsplusnames, i); - if (_PyUnicode_EQ(name, key)) { - if (read) { - if (framelocalsproxy_getval(frame->f_frame, co, i) != NULL) { - return i; - } - } else { - if (!(_PyLocals_GetKind(co->co_localspluskinds, i) & CO_FAST_HIDDEN)) { - return i; - } + if (found) { + // This is an attempt to read an unset local variable or + // write to a variable that is hidden from regular write operations + return -1; + } + // This is unlikely, but we need to make sure. This means the key + // is not interned. + for (int i = 0; i < co->co_nlocalsplus; i++) { + PyObject *name = PyTuple_GET_ITEM(co->co_localsplusnames, i); + Py_hash_t name_hash = PyObject_Hash(name); + assert(name_hash != -1); // keys are exact unicode + if (name_hash != key_hash) { + continue; + } + int same = PyObject_RichCompareBool(name, key, Py_EQ); + if (same < 0) { + return -2; + } + if (same) { + if (read) { + if (framelocalsproxy_getval(frame->f_frame, co, i) != NULL) { + return i; + } + } else { + if (!(_PyLocals_GetKind(co->co_localspluskinds, i) & CO_FAST_HIDDEN)) { + return i; } } } @@ -109,13 +126,14 @@ framelocalsproxy_getitem(PyObject *self, PyObject *key) PyFrameObject* frame = ((PyFrameLocalsProxyObject*)self)->frame; PyCodeObject* co = _PyFrame_GetCode(frame->f_frame); - if (PyUnicode_CheckExact(key)) { - int i = framelocalsproxy_getkeyindex(frame, key, true); - if (i >= 0) { - PyObject *value = framelocalsproxy_getval(frame->f_frame, co, i); - assert(value != NULL); - return Py_NewRef(value); - } + int i = framelocalsproxy_getkeyindex(frame, key, true); + if (i == -2) { + return NULL; + } + if (i >= 0) { + PyObject *value = framelocalsproxy_getval(frame->f_frame, co, i); + assert(value != NULL); + return Py_NewRef(value); } // Okay not in the fast locals, try extra locals @@ -145,37 +163,38 @@ framelocalsproxy_setitem(PyObject *self, PyObject *key, PyObject *value) return -1; } - if (PyUnicode_CheckExact(key)) { - int i = framelocalsproxy_getkeyindex(frame, key, false); - if (i >= 0) { - _Py_Executors_InvalidateDependency(PyInterpreterState_Get(), co, 1); - - _PyLocals_Kind kind = _PyLocals_GetKind(co->co_localspluskinds, i); - _PyStackRef oldvalue = fast[i]; - PyObject *cell = NULL; - if (kind == CO_FAST_FREE) { - // The cell was set when the frame was created from - // the function's closure. - assert(oldvalue.bits != 0 && PyCell_Check(PyStackRef_AsPyObjectBorrow(oldvalue))); - cell = PyStackRef_AsPyObjectBorrow(oldvalue); - } else if (kind & CO_FAST_CELL && oldvalue.bits != 0) { - PyObject *as_obj = PyStackRef_AsPyObjectBorrow(oldvalue); - if (PyCell_Check(as_obj)) { - cell = as_obj; - } + int i = framelocalsproxy_getkeyindex(frame, key, false); + if (i == -2) { + return -1; + } + if (i >= 0) { + _Py_Executors_InvalidateDependency(PyInterpreterState_Get(), co, 1); + + _PyLocals_Kind kind = _PyLocals_GetKind(co->co_localspluskinds, i); + _PyStackRef oldvalue = fast[i]; + PyObject *cell = NULL; + if (kind == CO_FAST_FREE) { + // The cell was set when the frame was created from + // the function's closure. + assert(oldvalue.bits != 0 && PyCell_Check(PyStackRef_AsPyObjectBorrow(oldvalue))); + cell = PyStackRef_AsPyObjectBorrow(oldvalue); + } else if (kind & CO_FAST_CELL && oldvalue.bits != 0) { + PyObject *as_obj = PyStackRef_AsPyObjectBorrow(oldvalue); + if (PyCell_Check(as_obj)) { + cell = as_obj; } - if (cell != NULL) { - PyObject *oldvalue_o = PyCell_GET(cell); - if (value != oldvalue_o) { - PyCell_SET(cell, Py_XNewRef(value)); - Py_XDECREF(oldvalue_o); - } - } else if (value != PyStackRef_AsPyObjectBorrow(oldvalue)) { - PyStackRef_XCLOSE(fast[i]); - fast[i] = PyStackRef_FromPyObjectNew(value); + } + if (cell != NULL) { + PyObject *oldvalue_o = PyCell_GET(cell); + if (value != oldvalue_o) { + PyCell_SET(cell, Py_XNewRef(value)); + Py_XDECREF(oldvalue_o); } - return 0; + } else if (value != PyStackRef_AsPyObjectBorrow(oldvalue)) { + PyStackRef_XCLOSE(fast[i]); + fast[i] = PyStackRef_FromPyObjectNew(value); } + return 0; } // Okay not in the fast locals, try extra locals @@ -545,11 +564,12 @@ framelocalsproxy_contains(PyObject *self, PyObject *key) { PyFrameObject *frame = ((PyFrameLocalsProxyObject*)self)->frame; - if (PyUnicode_CheckExact(key)) { - int i = framelocalsproxy_getkeyindex(frame, key, true); - if (i >= 0) { - return 1; - } + int i = framelocalsproxy_getkeyindex(frame, key, true); + if (i == -2) { + return -1; + } + if (i >= 0) { + return 1; } PyObject *extra = ((PyFrameObject*)frame)->f_extra_locals;