Add a __pybind11_enum__ capsule to native enums containing some information needed for interop with them, plus a destructor that unregisters the enum when it's destroyed

Author: oremanj

include/pybind11/detail/internals.h

@@ -431,6 +431,15 @@ struct type_info {
     bool module_local : 1;
 };
 
+/// Information stored in a capsule on py::native_enum() types.
+struct native_enum_info {
+    const std::type_info *cpptype;
+    uint32_t size_bytes;
+    bool is_signed;
+
+    static const char *attribute_name() { return "__pybind11_enum__"; }
+};
+
 #define PYBIND11_ABI_TAG                                                                          \
     "v" PYBIND11_TOSTRING(PYBIND11_INTERNALS_VERSION)                                             \
         PYBIND11_COMPILER_TYPE_LEADING_UNDERSCORE PYBIND11_PLATFORM_ABI_ID

include/pybind11/detail/native_enum_data.h

@@ -24,15 +24,25 @@ native_enum_missing_finalize_error_message(const std::string &enum_name_encoded)
 
 class native_enum_data {
 public:
-    native_enum_data(const object &parent_scope,
+    native_enum_data(handle parent_scope,
                      const char *enum_name,
                      const char *native_type_name,
                      const char *class_doc,
-                     const std::type_index &enum_type_index)
+                     const native_enum_info &enum_info_)
         : enum_name_encoded{enum_name}, native_type_name_encoded{native_type_name},
-          enum_type_index{enum_type_index}, parent_scope(parent_scope), enum_name{enum_name},
+          enum_type_index{*enum_info_.cpptype}, parent_scope(parent_scope), enum_name{enum_name},
           native_type_name{native_type_name}, class_doc(class_doc), export_values_flag{false},
-          finalize_needed{false} {}
+          finalize_needed{false} {
+        enum_info = capsule(new native_enum_info{enum_info_},
+                            native_enum_info::attribute_name(),
+                            +[](void *enum_info_) {
+                                auto *info = (native_enum_info *) enum_info_;
+                                with_internals([&](internals &internals) {
+                                    internals.native_enum_type_map.erase(*info->cpptype);
+                                });
+                                delete info;
+                            });
+    }
 
     void finalize();
 
@@ -67,10 +77,11 @@ class native_enum_data {
     std::type_index enum_type_index;
 
 private:
-    object parent_scope;
+    handle parent_scope;
     str enum_name;
     str native_type_name;
     std::string class_doc;
+    capsule enum_info;
 
 protected:
     list members;
@@ -81,12 +92,6 @@ class native_enum_data {
     bool finalize_needed : 1;
 };
 
-inline void global_internals_native_enum_type_map_set_item(const std::type_index &enum_type_index,
-                                                           PyObject *py_enum) {
-    with_internals(
-        [&](internals &internals) { internals.native_enum_type_map[enum_type_index] = py_enum; });
-}
-
 inline handle
 global_internals_native_enum_type_map_get_item(const std::type_index &enum_type_index) {
     return with_internals([&](internals &internals) {
@@ -202,7 +207,11 @@ inline void native_enum_data::finalize() {
     for (auto doc : member_docs) {
         py_enum[doc[int_(0)]].attr("__doc__") = doc[int_(1)];
     }
-    global_internals_native_enum_type_map_set_item(enum_type_index, py_enum.release().ptr());
+
+    py_enum.attr(native_enum_info::attribute_name()) = enum_info;
+    with_internals([&](internals &internals) {
+        internals.native_enum_type_map[enum_type_index] = py_enum.ptr();
+    });
 }
 
 PYBIND11_NAMESPACE_END(detail)