Support frozenset, tuple as dict keys

ecatmur · ecatmur · commit a4b56286cee9 · 2022-04-20T01:20:30.000+01:00
#3836 Add frozenset as a pybind11 type. Add freeze() function converting set to frozenset and list to tuple; use it in std::set and std::map casters. Add tests.
diff --git a/include/pybind11/cast.h b/include/pybind11/cast.h
@@ -1617,6 +1617,25 @@ object object_api<Derived>::call(Args &&...args) const {
     return operator()<policy>(std::forward<Args>(args)...);
 }
 
+// Convert list -> tuple and set -> frozenset for use as keys in dict, set etc.
+// https://mail.python.org/pipermail/python-dev/2005-October/057586.html
+inline object freeze(object &&obj) {
+    if (isinstance<list>(obj)) {
+        return tuple(std::move(obj));
+    } else if (isinstance<set>(obj)) {
+        return frozenset(std::move(obj));
+    } else {
+        return std::move(obj);
+    }
+}
+
+template <typename Caster, typename SFINAE = decltype(Caster::frozen_name)>
+constexpr inline auto get_frozen_name_impl(int) { return Caster::frozen_name; }
+template <typename Caster>
+constexpr inline auto get_frozen_name_impl(long) { return Caster::name; }
+template <typename Caster>
+constexpr inline auto get_frozen_name() { return get_frozen_name_impl<Caster>(0); }
+
 PYBIND11_NAMESPACE_END(detail)
 
 template <typename T>
diff --git a/include/pybind11/pytypes.h b/include/pybind11/pytypes.h
@@ -1784,24 +1784,41 @@ class kwargs : public dict {
     PYBIND11_OBJECT_DEFAULT(kwargs, dict, PyDict_Check)
 };
 
-class set : public object {
+class set_base : public object {
+protected:
+    PYBIND11_OBJECT(set_base, object, PyAnySet_Check)
+
+public:
+    size_t size() const { return (size_t) PySet_Size(m_ptr); }
+    bool empty() const { return size() == 0; }
+    template <typename T>
+    bool contains(T &&val) const {
+        return PySet_Contains(m_ptr, detail::object_or_cast(std::forward<T>(val)).ptr()) == 1;
+    }
+};
+
+class set : public set_base {
 public:
-    PYBIND11_OBJECT_CVT(set, object, PySet_Check, PySet_New)
-    set() : object(PySet_New(nullptr), stolen_t{}) {
+    PYBIND11_OBJECT_CVT(set, set_base, PySet_Check, PySet_New)
+    set() : set_base(PySet_New(nullptr), stolen_t{}) {
         if (!m_ptr) {
             pybind11_fail("Could not allocate set object!");
         }
     }
-    size_t size() const { return (size_t) PySet_Size(m_ptr); }
-    bool empty() const { return size() == 0; }
     template <typename T>
     bool add(T &&val) /* py-non-const */ {
         return PySet_Add(m_ptr, detail::object_or_cast(std::forward<T>(val)).ptr()) == 0;
     }
     void clear() /* py-non-const */ { PySet_Clear(m_ptr); }
-    template <typename T>
-    bool contains(T &&val) const {
-        return PySet_Contains(m_ptr, detail::object_or_cast(std::forward<T>(val)).ptr()) == 1;
+};
+
+class frozenset : public set_base {
+public:
+    PYBIND11_OBJECT_CVT(frozenset, set_base, PyFrozenSet_Check, PyFrozenSet_New)
+    frozenset() : set_base(PyFrozenSet_New(nullptr), stolen_t{}) {
+        if (!m_ptr) {
+            pybind11_fail("Could not allocate frozenset object!");
+        }
     }
 };
 
diff --git a/include/pybind11/stl.h b/include/pybind11/stl.h
@@ -55,10 +55,10 @@ struct set_caster {
     using key_conv = make_caster<Key>;
 
     bool load(handle src, bool convert) {
-        if (!isinstance<pybind11::set>(src)) {
+        if (!isinstance<set_base>(src)) {
             return false;
         }
-        auto s = reinterpret_borrow<pybind11::set>(src);
+        auto s = reinterpret_borrow<set_base>(src);
         value.clear();
         for (auto entry : s) {
             key_conv conv;
@@ -79,14 +79,15 @@ struct set_caster {
         for (auto &&value : src) {
             auto value_ = reinterpret_steal<object>(
                 key_conv::cast(forward_like<T>(value), policy, parent));
-            if (!value_ || !s.add(std::move(value_))) {
+            if (!value_ || !s.add(freeze(std::move(value_)))) {
                 return handle();
             }
         }
         return s.release();
     }
 
-    PYBIND11_TYPE_CASTER(type, const_name("Set[") + key_conv::name + const_name("]"));
+    PYBIND11_TYPE_CASTER(type, const_name("Set[") + get_frozen_name<key_conv>() + const_name("]"));
+    static constexpr auto frozen_name = const_name("FrozenSet[") + get_frozen_name<key_conv>() + const_name("]");
 };
 
 template <typename Type, typename Key, typename Value>
@@ -128,14 +129,14 @@ struct map_caster {
             if (!key || !value) {
                 return handle();
             }
-            d[key] = value;
+            d[freeze(std::move(key))] = std::move(value);
         }
         return d.release();
     }
 
     PYBIND11_TYPE_CASTER(Type,
-                         const_name("Dict[") + key_conv::name + const_name(", ") + value_conv::name
-                             + const_name("]"));
+                         const_name("Dict[") + get_frozen_name<key_conv>() + const_name(", ")
+                             + value_conv::name + const_name("]"));
 };
 
 template <typename Type, typename Value>
@@ -188,6 +189,7 @@ struct list_caster {
     }
 
     PYBIND11_TYPE_CASTER(Type, const_name("List[") + value_conv::name + const_name("]"));
+    static constexpr auto frozen_name = const_name("Tuple[") + value_conv::name + const_name(", ...]");
 };
 
 template <typename Type, typename Alloc>
@@ -257,6 +259,11 @@ struct array_caster {
                                                      const_name("[") + const_name<Size>()
                                                          + const_name("]"))
                              + const_name("]"));
+    static constexpr auto frozen_name = const_name("Tuple[") + value_conv::name
+                             + const_name<Resizable>(const_name(", ..."),
+                                                     const_name("[") + const_name<Size>()
+                                                         + const_name("]"))
+                             + const_name("]");
 };
 
 template <typename Type, size_t Size>
diff --git a/tests/test_stl.cpp b/tests/test_stl.cpp
@@ -248,6 +248,22 @@ TEST_SUBMODULE(stl, m) {
         return v;
     });
 
+    // test_frozen_key
+    m.def("cast_set_map", []() {
+        return std::map<std::set<std::string>, std::string>{{{"key1", "key2"}, "value"}};
+    });
+    m.def("load_set_map", [](const std::map<std::set<std::string>, std::string> &map) {
+        return map.at({"key1", "key2"}) == "value" && map.at({"key3"}) == "value2";
+    });
+    m.def("cast_set_set", []() { return std::set<std::set<std::string>>{{"key1", "key2"}}; });
+    m.def("load_set_set", [](const std::set<std::set<std::string>> &set) {
+        return (set.count({"key1", "key2"}) != 0u) && (set.count({"key3"}) != 0u);
+    });
+    m.def("cast_vector_set", []() { return std::set<std::vector<int>>{{1, 2}}; });
+    m.def("load_vector_set", [](const std::set<std::vector<int>> &set) {
+        return (set.count({1, 2}) != 0u) && (set.count({3}) != 0u);
+    });
+
     pybind11::enum_<EnumType>(m, "EnumType")
         .value("kSet", EnumType::kSet)
         .value("kUnset", EnumType::kUnset);
diff --git a/tests/test_stl.py b/tests/test_stl.py
@@ -94,6 +94,38 @@ def test_recursive_casting():
     assert z[0].value == 7 and z[1].value == 42
 
 
+def test_frozen_key(doc):
+    """Test that we special-case C++ key types to Python immutable containers, e.g.:
+    std::map<std::set<K>, V> <-> dict[frozenset[K], V]
+    std::set<std::set<T>> <-> set[frozenset[T]]
+    std::set<std::vector<T>> <-> set[tuple[T, ...]]
+    """
+    s = m.cast_set_map()
+    assert s == {frozenset({"key1", "key2"}): "value"}
+    s[frozenset({"key3"})] = "value2"
+    assert m.load_set_map(s)
+    assert doc(m.cast_set_map) == "cast_set_map() -> Dict[FrozenSet[str], str]"
+    assert (
+        doc(m.load_set_map) == "load_set_map(arg0: Dict[FrozenSet[str], str]) -> bool"
+    )
+
+    s = m.cast_set_set()
+    assert s == {frozenset({"key1", "key2"})}
+    s.add(frozenset({"key3"}))
+    assert m.load_set_set(s)
+    assert doc(m.cast_set_set) == "cast_set_set() -> Set[FrozenSet[str]]"
+    assert doc(m.load_set_set) == "load_set_set(arg0: Set[FrozenSet[str]]) -> bool"
+
+    s = m.cast_vector_set()
+    assert s == {(1, 2)}
+    s.add((3,))
+    assert m.load_vector_set(s)
+    assert doc(m.cast_vector_set) == "cast_vector_set() -> Set[Tuple[int, ...]]"
+    assert (
+        doc(m.load_vector_set) == "load_vector_set(arg0: Set[Tuple[int, ...]]) -> bool"
+    )
+
+
 def test_move_out_container():
     """Properties use the `reference_internal` policy by default. If the underlying function
     returns an rvalue, the policy is automatically changed to `move` to avoid referencing
