Support frozenset, tuple as dict keys

include/pybind11/cast.h

@@ -79,6 +79,9 @@ class type_caster<std::reference_wrapper<type>> {
     explicit operator std::reference_wrapper<type>() { return cast_op<type &>(subcaster); }
 };
 
+template <typename type>
+class type_caster<const type> : public type_caster<type> {};
+
 #define PYBIND11_TYPE_CASTER(type, py_name)                                                       \
 protected:                                                                                        \
     type value;                                                                                   \

include/pybind11/pytypes.h

@@ -1784,14 +1784,11 @@ 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:
-    PYBIND11_OBJECT_CVT(set, object, PySet_Check, PySet_New)
-    set() : object(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>
@@ -1805,6 +1802,26 @@ class set : public object {
     }
 };
 
+class set : public set_base {
+public:
+    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!");
+        }
+    }
+};
+
+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!");
+        }
+    }
+};
+
 class function : public object {
 public:
     PYBIND11_OBJECT_DEFAULT(function, object, PyCallable_Check)

include/pybind11/stl.h

@@ -49,23 +49,24 @@ forwarded_type<T, U> forward_like(U &&u) {
     return std::forward<detail::forwarded_type<T, U>>(std::forward<U>(u));
 }
 
-template <typename Type, typename Key>
+template <typename Type, typename Key, bool Const = false>
 struct set_caster {
     using type = Type;
-    using key_conv = make_caster<Key>;
+    using key_conv = type_caster<const intrinsic_t<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;
             if (!conv.load(entry, convert)) {
                 return false;
             }
-            value.insert(cast_op<Key &&>(std::move(conv)));
+            value.insert(
+                std::move(conv).operator typename key_conv::template cast_op_type<Key &&>());
         }
         return true;
     }
@@ -75,7 +76,7 @@ struct set_caster {
         if (!std::is_lvalue_reference<T>::value) {
             policy = return_value_policy_override<Key>::policy(policy);
         }
-        pybind11::set s;
+        typename std::conditional<Const, frozenset, set>::type s;
         for (auto &&value : src) {
             auto value_ = reinterpret_steal<object>(
                 key_conv::cast(forward_like<T>(value), policy, parent));
@@ -86,12 +87,14 @@ struct set_caster {
         return s.release();
     }
 
-    PYBIND11_TYPE_CASTER(type, const_name("Set[") + key_conv::name + const_name("]"));
+    PYBIND11_TYPE_CASTER(type,
+                         const_name<Const>("FrozenSet[", "Set[") + key_conv::name
+                             + const_name("]"));
 };
 
 template <typename Type, typename Key, typename Value>
 struct map_caster {
-    using key_conv = make_caster<Key>;
+    using key_conv = type_caster<const intrinsic_t<Key>>;
     using value_conv = make_caster<Value>;
 
     bool load(handle src, bool convert) {
@@ -106,7 +109,9 @@ struct map_caster {
             if (!kconv.load(it.first.ptr(), convert) || !vconv.load(it.second.ptr(), convert)) {
                 return false;
             }
-            value.emplace(cast_op<Key &&>(std::move(kconv)), cast_op<Value &&>(std::move(vconv)));
+            value.emplace(
+                std::move(kconv).operator typename key_conv::template cast_op_type<Key &&>(),
+                cast_op<Value &&>(std::move(vconv)));
         }
         return true;
     }
@@ -138,7 +143,7 @@ struct map_caster {
                              + const_name("]"));
 };
 
-template <typename Type, typename Value>
+template <typename Type, typename Value, bool Const = false>
 struct list_caster {
     using value_conv = make_caster<Value>;
 
@@ -174,32 +179,50 @@ struct list_caster {
         if (!std::is_lvalue_reference<T>::value) {
             policy = return_value_policy_override<Value>::policy(policy);
         }
-        list l(src.size());
+        conditional_t<Const, tuple, list> l(src.size());
         ssize_t index = 0;
         for (auto &&value : src) {
             auto value_ = reinterpret_steal<object>(
                 value_conv::cast(forward_like<T>(value), policy, parent));
             if (!value_) {
                 return handle();
             }
-            PyList_SET_ITEM(l.ptr(), index++, value_.release().ptr()); // steals a reference
+            if (Const) {
+                PyTuple_SET_ITEM(l.ptr(), index++, value_.release().ptr()); // steals a reference
+            } else {
+                PyList_SET_ITEM(l.ptr(), index++, value_.release().ptr()); // steals a reference
+            }
         }
         return l.release();
     }
 
-    PYBIND11_TYPE_CASTER(Type, const_name("List[") + value_conv::name + const_name("]"));
+    PYBIND11_TYPE_CASTER(Type,
+                         const_name<Const>("Tuple[", "List[") + value_conv::name
+                             + const_name<Const>(", ...]", "]"));
 };
 
 template <typename Type, typename Alloc>
 struct type_caster<std::vector<Type, Alloc>> : list_caster<std::vector<Type, Alloc>, Type> {};
 
+template <typename Type, typename Alloc>
+struct type_caster<const std::vector<Type, Alloc>>
+    : list_caster<std::vector<Type, Alloc>, Type, true> {};
+
 template <typename Type, typename Alloc>
 struct type_caster<std::deque<Type, Alloc>> : list_caster<std::deque<Type, Alloc>, Type> {};
 
+template <typename Type, typename Alloc>
+struct type_caster<const std::deque<Type, Alloc>>
+    : list_caster<std::deque<Type, Alloc>, Type, true> {};
+
 template <typename Type, typename Alloc>
 struct type_caster<std::list<Type, Alloc>> : list_caster<std::list<Type, Alloc>, Type> {};
 
-template <typename ArrayType, typename Value, bool Resizable, size_t Size = 0>
+template <typename Type, typename Alloc>
+struct type_caster<const std::list<Type, Alloc>>
+    : list_caster<std::list<Type, Alloc>, Type, true> {};
+
+template <typename ArrayType, typename Value, bool Resizable, size_t Size = 0, bool Const = false>
 struct array_caster {
     using value_conv = make_caster<Value>;
 
@@ -238,42 +261,58 @@ struct array_caster {
 
     template <typename T>
     static handle cast(T &&src, return_value_policy policy, handle parent) {
-        list l(src.size());
+        conditional_t<Const, tuple, list> l(src.size());
         ssize_t index = 0;
         for (auto &&value : src) {
             auto value_ = reinterpret_steal<object>(
                 value_conv::cast(forward_like<T>(value), policy, parent));
             if (!value_) {
                 return handle();
             }
-            PyList_SET_ITEM(l.ptr(), index++, value_.release().ptr()); // steals a reference
+            if (Const) {
+                PyTuple_SET_ITEM(l.ptr(), index++, value_.release().ptr()); // steals a reference
+            } else {
+                PyList_SET_ITEM(l.ptr(), index++, value_.release().ptr()); // steals a reference
+            }
         }
         return l.release();
     }
 
     PYBIND11_TYPE_CASTER(ArrayType,
-                         const_name("List[") + value_conv::name
+                         const_name<Const>("Tuple[", "List[") + value_conv::name
                              + const_name<Resizable>(const_name(""),
                                                      const_name("[") + const_name<Size>()
                                                          + const_name("]"))
-                             + const_name("]"));
+                             + const_name<Const>(", ...]", "]"));
 };
 
 template <typename Type, size_t Size>
 struct type_caster<std::array<Type, Size>>
     : array_caster<std::array<Type, Size>, Type, false, Size> {};
 
+template <typename Type, size_t Size>
+struct type_caster<const std::array<Type, Size>>
+    : array_caster<std::array<Type, Size>, Type, false, Size, true> {};
+
 template <typename Type>
 struct type_caster<std::valarray<Type>> : array_caster<std::valarray<Type>, Type, true> {};
 
 template <typename Key, typename Compare, typename Alloc>
 struct type_caster<std::set<Key, Compare, Alloc>>
     : set_caster<std::set<Key, Compare, Alloc>, Key> {};
 
+template <typename Key, typename Compare, typename Alloc>
+struct type_caster<const std::set<Key, Compare, Alloc>>
+    : set_caster<std::set<Key, Compare, Alloc>, Key, true> {};
+
 template <typename Key, typename Hash, typename Equal, typename Alloc>
 struct type_caster<std::unordered_set<Key, Hash, Equal, Alloc>>
     : set_caster<std::unordered_set<Key, Hash, Equal, Alloc>, Key> {};
 
+template <typename Key, typename Hash, typename Equal, typename Alloc>
+struct type_caster<const std::unordered_set<Key, Hash, Equal, Alloc>>
+    : set_caster<std::unordered_set<Key, Hash, Equal, Alloc>, Key, true> {};
+
 template <typename Key, typename Value, typename Compare, typename Alloc>
 struct type_caster<std::map<Key, Value, Compare, Alloc>>
     : map_caster<std::map<Key, Value, Compare, Alloc>, Key, Value> {};

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);

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