Support non-string keys

markshannon · May 4, 2024 · Feb 7, 2024 · Feb 8, 2024 · Feb 8, 2024 · Feb 8, 2024
commit 8dc46642fa98feb1ff0b14c3da392b52e90ec971
diff --git a/Lib/test/test_frame.py b/Lib/test/test_frame.py
@@ -331,6 +331,9 @@ def test_as_dict(self):
         self.assertEqual(d.setdefault('new', 5), 5)
         self.assertEqual(d['new'], 5)
 
+        with self.assertRaises(KeyError):
+            d['non_exist']
+
     def test_as_number(self):
         x = 1
         y = 2
@@ -345,6 +348,11 @@ def test_as_number(self):
         with self.assertRaises(TypeError):
             _ = d | [3]
 
+    def test_non_string_key(self):
+        d = sys._getframe().f_locals
+        d[1] = 2
+        self.assertEqual(d[1], 2)
+
     def test_write_with_hidden(self):
         def f():
             f_locals = [sys._getframe().f_locals for b in [0]][0]
@@ -391,12 +399,6 @@ def test_unsupport(self):
         with self.assertRaises(TypeError):
             copy.deepcopy(d)
 
-        with self.assertRaises(TypeError):
-            d.get(1)
-
-        with self.assertRaises(TypeError):
-            d.setdefault(1, 'x')
-
 
 class TestFrameCApi(unittest.TestCase):
     def test_basic(self):

@@ -59,6 +59,9 @@ framelocalsproxy_getkeyindex(PyFrameObject *frame, PyObject* key, bool read)
      *   - if read == true, returns the index if the value is not NULL
      *   - if read == false, returns the index if the value is not hidden
      */
+
+    assert(PyUnicode_CheckExact(key));
+
     PyCodeObject *co = PyFrame_GetCode(frame);
     int found_key = false;
 
@@ -388,11 +391,13 @@ framelocalsproxy_getitem(PyObject *self, PyObject *key)
     PyFrameObject* frame = ((PyFrameLocalsProxyObject*)self)->frame;
     PyCodeObject* co = PyFrame_GetCode(frame);
 
-    int i = framelocalsproxy_getkeyindex(frame, key, true);
-    if (i >= 0) {
-        PyObject *value = framelocalsproxy_getval(frame, co, i);
-        assert(value != NULL);
-        return Py_NewRef(value);
+    if (PyUnicode_CheckExact(key)) {
+        int i = framelocalsproxy_getkeyindex(frame, key, true);
+        if (i >= 0) {
+            PyObject *value = framelocalsproxy_getval(frame, co, i);
+            assert(value != NULL);
+            return Py_NewRef(value);
+        }
     }
 
     // Okay not in the fast locals, try extra locals
@@ -422,34 +427,35 @@ framelocalsproxy_setitem(PyObject *self, PyObject *key, PyObject *value)
         return -1;
     }
 
-    int i = framelocalsproxy_getkeyindex(frame, key, false);
-
-    if (i >= 0) {
-        _PyLocals_Kind kind = _PyLocals_GetKind(co->co_localspluskinds, i);
+    if (PyUnicode_CheckExact(key)) {
+        int i = framelocalsproxy_getkeyindex(frame, key, false);
+        if (i >= 0) {
+            _PyLocals_Kind kind = _PyLocals_GetKind(co->co_localspluskinds, i);
 
-        PyObject *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 != NULL && PyCell_Check(oldvalue));
-            cell = oldvalue;
-        } else if (kind & CO_FAST_CELL && oldvalue != NULL) {
-            if (PyCell_Check(oldvalue)) {
+            PyObject *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 != NULL && PyCell_Check(oldvalue));
                 cell = oldvalue;
+            } else if (kind & CO_FAST_CELL && oldvalue != NULL) {
+                if (PyCell_Check(oldvalue)) {
+                    cell = oldvalue;
+                }
             }
-        }
-        if (cell != NULL) {
-            oldvalue = PyCell_GET(cell);
-            if (value != oldvalue) {
-                PyCell_SET(cell, Py_XNewRef(value));
-                Py_XDECREF(oldvalue);
+            if (cell != NULL) {
+                oldvalue = PyCell_GET(cell);
+                if (value != oldvalue) {
+                    PyCell_SET(cell, Py_XNewRef(value));
+                    Py_XDECREF(oldvalue);
+                }
+            } else if (value != oldvalue) {
+                Py_XSETREF(fast[i], Py_NewRef(value));
             }
-        } else if (value != oldvalue) {
-            Py_XSETREF(fast[i], Py_NewRef(value));
+            Py_XDECREF(value);
+            return 0;
         }
-        Py_XDECREF(value);
-        return 0;
     }
 
     // Okay not in the fast locals, try extra locals
@@ -473,21 +479,21 @@ static PyObject*
 framelocalsproxy_contains(PyObject *self, PyObject *key)
 {
     PyFrameObject *frame = ((PyFrameLocalsProxyObject*)self)->frame;
-    int i = framelocalsproxy_getkeyindex(frame, key, true);
-
-    if (i >= 0) {
-        PyObject *value = framelocalsproxy_getval(frame, PyFrame_GetCode(frame), i);
-        assert(value != NULL);
-        Py_RETURN_TRUE;
-    } else {
-        PyObject *extra = ((PyFrameObject*)frame)->f_extra_locals;
-        if (extra != NULL) {
-            int result = PyDict_Contains(extra, key);
-            if (result < 0) {
-                return NULL;
-            } else if (result > 0) {
-                Py_RETURN_TRUE;
-            }
+
+    if (PyUnicode_CheckExact(key)) {
+        int i = framelocalsproxy_getkeyindex(frame, key, true);
+        if (i >= 0) {
+            Py_RETURN_TRUE;
+        }
+    }
+
+    PyObject *extra = ((PyFrameObject*)frame)->f_extra_locals;
+    if (extra != NULL) {
+        int result = PyDict_Contains(extra, key);
+        if (result < 0) {
+            return NULL;
+        } else if (result > 0) {
+            Py_RETURN_TRUE;
         }
     }
 
@@ -520,12 +526,7 @@ framelocalsproxy_get(PyObject* self, PyObject *const *args, Py_ssize_t nargs)
         default_value = args[1];
     }
 
-    if (!PyUnicode_CheckExact(key)) {
-        PyErr_SetString(PyExc_TypeError, "key must be a string");
-        return NULL;
-    }
-
-    PyObject *result = framelocalsproxy_getitem(self, args[0]);
+    PyObject *result = framelocalsproxy_getitem(self, key);
 
     if (result == NULL) {
         if (PyErr_ExceptionMatches(PyExc_KeyError)) {
@@ -553,11 +554,6 @@ framelocalsproxy_setdefault(PyObject* self, PyObject *const *args, Py_ssize_t na
         default_value = args[1];
     }
 
-    if (!PyUnicode_CheckExact(key)) {
-        PyErr_SetString(PyExc_TypeError, "key must be a string");
-        return NULL;
-    }
-
     PyObject *result = framelocalsproxy_getitem(self, key);
 
     if (result == NULL) {