gh-120906: Support arbitrary hashable keys in FrameLocalsProxy (GH-122309)

encukou · ncoghlan · web-flow · commit 5912487938ac · 2024-07-30T22:11:00.000Z
Co-authored-by: Alyssa Coghlan &lt;ncoghlan@gmail.com&gt;
diff --git 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
@@ -0,0 +1 @@
+:attr:`frame.f_locals` now supports arbitrary hashable objects as keys.
diff --git 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;

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1 @@`
	`1`	+:attr:`frame.f_locals` now supports arbitrary hashable objects as keys.