From 56e69a20a5b4ab6d31a262dde5ef723a322033a5 Mon Sep 17 00:00:00 2001
From: Paul-Edouard Sarlin <15985472+sarlinpe@users.noreply.github.com>
Date: Tue, 8 Oct 2024 19:49:35 +0200
Subject: [PATCH] Print key in KeyError in map.__getitem__/__delitem__ (#5397)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* Print key in map.getitem/delitem KeyError

* Add tests

* Fix tests

* Make robust

* Make clang-tidy happy

* Return a Python str

* Show beginning and end of the message

* Avoid implicit conversion

* Split out `format_message_key_error_key_object()` to reduce amount of templated code.

* Use `"✄✄✄"` instead of `"..."`

Also rename variable to `cut_length`, to not get into even/odd issues with the meaning of "half".

---------

Co-authored-by: Ralf W. Grosse-Kunstleve <rgrossekunst@nvidia.com>
---
 include/pybind11/stl_bind.h | 40 +++++++++++++++++++++++++++++++++++--
 tests/test_stl_binders.py   | 19 ++++++++++++++++++
 2 files changed, 57 insertions(+), 2 deletions(-)

diff --git a/include/pybind11/stl_bind.h b/include/pybind11/stl_bind.h
index fcb48dea3..af3a47f39 100644
--- a/include/pybind11/stl_bind.h
+++ b/include/pybind11/stl_bind.h
@@ -694,6 +694,40 @@ struct ItemsViewImpl : public detail::items_view {
     Map &map;
 };
 
+inline str format_message_key_error_key_object(handle py_key) {
+    str message = "pybind11::bind_map key";
+    if (!py_key) {
+        return message;
+    }
+    try {
+        message = str(py_key);
+    } catch (const std::exception &) {
+        try {
+            message = repr(py_key);
+        } catch (const std::exception &) {
+            return message;
+        }
+    }
+    const ssize_t cut_length = 100;
+    if (len(message) > 2 * cut_length + 3) {
+        return str(message[slice(0, cut_length, 1)]) + str("✄✄✄")
+               + str(message[slice(-cut_length, static_cast<ssize_t>(len(message)), 1)]);
+    }
+    return message;
+}
+
+template <typename KeyType>
+str format_message_key_error(const KeyType &key) {
+    object py_key;
+    try {
+        py_key = cast(key);
+    } catch (const std::exception &) {
+        do { // Trick to avoid "empty catch" warning/error.
+        } while (false);
+    }
+    return format_message_key_error_key_object(py_key);
+}
+
 PYBIND11_NAMESPACE_END(detail)
 
 template <typename Map, typename holder_type = std::unique_ptr<Map>, typename... Args>
@@ -785,7 +819,8 @@ class_<Map, holder_type> bind_map(handle scope, const std::string &name, Args &&
         [](Map &m, const KeyType &k) -> MappedType & {
             auto it = m.find(k);
             if (it == m.end()) {
-                throw key_error();
+                set_error(PyExc_KeyError, detail::format_message_key_error(k));
+                throw error_already_set();
             }
             return it->second;
         },
@@ -808,7 +843,8 @@ class_<Map, holder_type> bind_map(handle scope, const std::string &name, Args &&
     cl.def("__delitem__", [](Map &m, const KeyType &k) {
         auto it = m.find(k);
         if (it == m.end()) {
-            throw key_error();
+            set_error(PyExc_KeyError, detail::format_message_key_error(k));
+            throw error_already_set();
         }
         m.erase(it);
     });
diff --git a/tests/test_stl_binders.py b/tests/test_stl_binders.py
index 09e784e2e..9856ba462 100644
--- a/tests/test_stl_binders.py
+++ b/tests/test_stl_binders.py
@@ -302,6 +302,25 @@ def test_map_delitem():
     assert list(mm) == ["b"]
     assert list(mm.items()) == [("b", 2.5)]
 
+    with pytest.raises(KeyError) as excinfo:
+        mm["a_long_key"]
+    assert "a_long_key" in str(excinfo.value)
+
+    with pytest.raises(KeyError) as excinfo:
+        del mm["a_long_key"]
+    assert "a_long_key" in str(excinfo.value)
+
+    cut_length = 100
+    k_very_long = "ab" * cut_length + "xyz"
+    with pytest.raises(KeyError) as excinfo:
+        mm[k_very_long]
+    assert k_very_long in str(excinfo.value)
+    k_very_long += "@"
+    with pytest.raises(KeyError) as excinfo:
+        mm[k_very_long]
+    k_repr = k_very_long[:cut_length] + "✄✄✄" + k_very_long[-cut_length:]
+    assert k_repr in str(excinfo.value)
+
     um = m.UnorderedMapStringDouble()
     um["ua"] = 1.1
     um["ub"] = 2.6