native_enum: add capsule containing enum information and cleanup logic (#5871)

* native_enum: add capsule containing enum information and cleanup logic

* style: pre-commit fixes

* Updates from code review

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: oremanj

include/pybind11/detail/internals.h

@@ -382,6 +382,24 @@ struct type_info {
     bool module_local : 1;
 };
 
+/// Information stored in a capsule on py::native_enum() types. Since we don't
+/// create a type_info record for native enums, we must store here any
+/// information we will need about the enum at runtime.
+///
+/// If you make backward-incompatible changes to this structure, you must
+/// change the `attribute_name()` so that native enums from older version of
+/// pybind11 don't have their records reinterpreted. Better would be to keep
+/// the changes backward-compatible (i.e., only add new fields at the end)
+/// and detect/indicate their presence using the currently-unused `version`.
+struct native_enum_record {
+    const std::type_info *cpptype;
+    uint32_t size_bytes;
+    bool is_signed;
+    const uint8_t version = 1;
+
+    static const char *attribute_name() { return "__pybind11_native_enum__"; }
+};
+
 #define PYBIND11_INTERNALS_ID                                                                     \
     "__pybind11_internals_v" PYBIND11_TOSTRING(PYBIND11_INTERNALS_VERSION)                        \
         PYBIND11_COMPILER_TYPE_LEADING_UNDERSCORE PYBIND11_PLATFORM_ABI_ID "__"

include/pybind11/detail/native_enum_data.h

@@ -22,17 +22,36 @@ native_enum_missing_finalize_error_message(const std::string &enum_name_encoded)
     return "pybind11::native_enum<...>(\"" + enum_name_encoded + "\", ...): MISSING .finalize()";
 }
 
+// Internals for pybind11::native_enum; one native_enum_data object exists
+// inside each pybind11::native_enum and lives only for the duration of the
+// native_enum binding statement.
 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_record &enum_record_)
         : 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_record_.cpptype},
+          parent_scope(reinterpret_borrow<object>(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} {
+        // Create the enum record capsule. It will be installed on the enum
+        // type object during finalize(). Its destructor removes the enum
+        // mapping from our internals, so that we won't try to convert to an
+        // enum type that's been destroyed.
+        enum_record = capsule(
+            new native_enum_record{enum_record_},
+            native_enum_record::attribute_name(),
+            +[](void *record_) {
+                auto *record = static_cast<native_enum_record *>(record_);
+                with_internals([&](internals &internals) {
+                    internals.native_enum_type_map.erase(*record->cpptype);
+                });
+                delete record;
+            });
+    }
 
     void finalize();
 
@@ -71,6 +90,7 @@ class native_enum_data {
     str enum_name;
     str native_type_name;
     std::string class_doc;
+    capsule enum_record;
 
 protected:
     list members;
@@ -81,12 +101,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 +216,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_record::attribute_name()) = enum_record;
+    with_internals([&](internals &internals) {
+        internals.native_enum_type_map[enum_type_index] = py_enum.ptr();
+    });
 }
 
 PYBIND11_NAMESPACE_END(detail)

include/pybind11/native_enum.h

@@ -22,12 +22,12 @@ class native_enum : public detail::native_enum_data {
 public:
     using Underlying = typename std::underlying_type<EnumType>::type;
 
-    native_enum(const object &parent_scope,
+    native_enum(handle parent_scope,
                 const char *name,
                 const char *native_type_name,
                 const char *class_doc = "")
         : detail::native_enum_data(
-              parent_scope, name, native_type_name, class_doc, std::type_index(typeid(EnumType))) {
+              parent_scope, name, native_type_name, class_doc, make_record()) {
         if (detail::get_local_type_info(typeid(EnumType)) != nullptr
             || detail::get_global_type_info(typeid(EnumType)) != nullptr) {
             pybind11_fail(
@@ -62,6 +62,15 @@ class native_enum : public detail::native_enum_data {
 
     native_enum(const native_enum &) = delete;
     native_enum &operator=(const native_enum &) = delete;
+
+private:
+    static detail::native_enum_record make_record() {
+        detail::native_enum_record ret;
+        ret.cpptype = &typeid(EnumType);
+        ret.size_bytes = sizeof(EnumType);
+        ret.is_signed = std::is_signed<Underlying>::value;
+        return ret;
+    }
 };
 
 PYBIND11_NAMESPACE_END(PYBIND11_NAMESPACE)

tests/conftest.py

@@ -16,6 +16,7 @@
 import sysconfig
 import textwrap
 import traceback
+import weakref
 from typing import Callable
 
 import pytest
@@ -208,19 +209,54 @@ def pytest_assertrepr_compare(op, left, right):  # noqa: ARG001
     return None
 
 
+# Number of times we think repeatedly collecting garbage might do anything.
+# The only reason to do more than once is because finalizers executed during
+# one GC pass could create garbage that can't be collected until a future one.
+# This quickly produces diminishing returns, and GC passes can be slow, so this
+# value is a tradeoff between non-flakiness and fast tests. (It errs on the
+# side of non-flakiness; many uses of this idiom only do 3 passes.)
+num_gc_collect = 5
+
+
 def gc_collect():
-    """Run the garbage collector three times (needed when running
+    """Run the garbage collector several times (needed when running
     reference counting tests with PyPy)"""
-    gc.collect()
-    gc.collect()
-    gc.collect()
-    gc.collect()
-    gc.collect()
+    for _ in range(num_gc_collect):
+        gc.collect()
+
+
+def delattr_and_ensure_destroyed(*specs):
+    """For each of the given *specs* (a tuple of the form ``(scope, name)``),
+    perform ``delattr(scope, name)``, then do enough GC collections that the
+    deleted reference has actually caused the target to be destroyed. This is
+    typically used to test what happens when a type object is destroyed; if you
+    use it for that, you should be aware that extension types, or all types,
+    are immortal on some Python versions. See ``env.TYPES_ARE_IMMORTAL``.
+    """
+    wrs = []
+    for mod, name in specs:
+        wrs.append(weakref.ref(getattr(mod, name)))
+        delattr(mod, name)
+
+    for _ in range(num_gc_collect):
+        gc.collect()
+        if all(wr() is None for wr in wrs):
+            break
+    else:
+        # If this fires, most likely something is still holding a reference
+        # to the object you tried to destroy - for example, it's a type that
+        # still has some instances alive. Try setting a breakpoint here and
+        # examining `gc.get_referrers(wrs[0]())`. It's vaguely possible that
+        # num_gc_collect needs to be increased also.
+        pytest.fail(
+            f"Could not delete bindings such as {next(wr for wr in wrs if wr() is not None)!r}"
+        )
 
 
 def pytest_configure():
     pytest.suppress = contextlib.suppress
     pytest.gc_collect = gc_collect
+    pytest.delattr_and_ensure_destroyed = delattr_and_ensure_destroyed
 
 
 def pytest_report_header():

tests/env.py

@@ -24,6 +24,12 @@
 # Runtime state (what's actually happening now)
 sys_is_gil_enabled = getattr(sys, "_is_gil_enabled", lambda: True)
 
+TYPES_ARE_IMMORTAL = (
+    PYPY
+    or GRAALPY
+    or (CPYTHON and PY_GIL_DISABLED and (3, 13) <= sys.version_info < (3, 14))
+)
+
 
 def deprecated_call():
     """

tests/test_class_cross_module_use_after_one_module_dealloc.py

@@ -1,40 +1,16 @@
 from __future__ import annotations
 
-import gc
-import sys
-import sysconfig
 import types
-import weakref
 
 import pytest
 
 import env
 from pybind11_tests import class_cross_module_use_after_one_module_dealloc as m
 
-is_python_3_13_free_threaded = (
-    env.CPYTHON
-    and sysconfig.get_config_var("Py_GIL_DISABLED")
-    and (3, 13) <= sys.version_info < (3, 14)
-)
-
-
-def delattr_and_ensure_destroyed(*specs):
-    wrs = []
-    for mod, name in specs:
-        wrs.append(weakref.ref(getattr(mod, name)))
-        delattr(mod, name)
 
-    for _ in range(5):
-        gc.collect()
-        if all(wr() is None for wr in wrs):
-            break
-    else:
-        pytest.fail(
-            f"Could not delete bindings such as {next(wr for wr in wrs if wr() is not None)!r}"
-        )
-
-
-@pytest.mark.skipif("env.PYPY or env.GRAALPY or is_python_3_13_free_threaded")
+@pytest.mark.skipif(
+    env.TYPES_ARE_IMMORTAL, reason="can't GC type objects on this platform"
+)
 def test_cross_module_use_after_one_module_dealloc():
     # This is a regression test for a bug that occurred during development of
     # internals::registered_types_cpp_fast (see #5842). registered_types_cpp_fast maps
@@ -58,7 +34,7 @@ def test_cross_module_use_after_one_module_dealloc():
     cm.consume_cross_dso_class(instance)
 
     del instance
-    delattr_and_ensure_destroyed((module_scope, "CrossDSOClass"))
+    pytest.delattr_and_ensure_destroyed((module_scope, "CrossDSOClass"))
 
     # Make sure that CrossDSOClass gets allocated at a different address.
     m.register_unrelated_class(module_scope)

tests/test_native_enum.cpp

@@ -93,10 +93,14 @@ TEST_SUBMODULE(native_enum, m) {
         .value("blue", color::blue)
         .finalize();
 
-    py::native_enum<altitude>(m, "altitude", "enum.Enum")
-        .value("high", altitude::high)
-        .value("low", altitude::low)
-        .finalize();
+    m.def("bind_altitude", [](const py::module_ &mod) {
+        py::native_enum<altitude>(mod, "altitude", "enum.Enum")
+            .value("high", altitude::high)
+            .value("low", altitude::low)
+            .finalize();
+    });
+    m.def("is_high_altitude", [](altitude alt) { return alt == altitude::high; });
+    m.def("get_altitude", []() -> altitude { return altitude::high; });
 
     py::native_enum<flags_uchar>(m, "flags_uchar", "enum.Flag")
         .value("bit0", flags_uchar::bit0)

tests/test_native_enum.py

@@ -59,7 +59,6 @@
 ENUM_TYPES_AND_MEMBERS = (
     (m.smallenum, SMALLENUM_MEMBERS),
     (m.color, COLOR_MEMBERS),
-    (m.altitude, ALTITUDE_MEMBERS),
     (m.flags_uchar, FLAGS_UCHAR_MEMBERS),
     (m.flags_uint, FLAGS_UINT_MEMBERS),
     (m.export_values, EXPORT_VALUES_MEMBERS),
@@ -320,3 +319,59 @@ def test_native_enum_missing_finalize_failure():
     if not isinstance(m.native_enum_missing_finalize_failure, str):
         m.native_enum_missing_finalize_failure()
         pytest.fail("Process termination expected.")
+
+
+def test_unregister_native_enum_when_destroyed():
+    # For stability when running tests in parallel, this test should be the
+    # only one that touches `m.altitude` or calls `m.bind_altitude`.
+
+    def test_altitude_enum():
+        # Logic copied from test_enum_type / test_enum_members.
+        # We don't test altitude there to avoid possible clashes if
+        # parallelizing against other tests in this file, and we also
+        # don't want to hold any references to the enumerators that
+        # would prevent GCing the enum type below.
+        assert isinstance(m.altitude, enum.EnumMeta)
+        assert m.altitude.__module__ == m.__name__
+        for name, value in ALTITUDE_MEMBERS:
+            assert m.altitude[name].value == value
+
+    def test_altitude_binding():
+        assert m.is_high_altitude(m.altitude.high)
+        assert not m.is_high_altitude(m.altitude.low)
+        assert m.get_altitude() is m.altitude.high
+        with pytest.raises(TypeError, match="incompatible function arguments"):
+            m.is_high_altitude("oops")
+
+    m.bind_altitude(m)
+    test_altitude_enum()
+    test_altitude_binding()
+
+    if env.TYPES_ARE_IMMORTAL:
+        pytest.skip("can't GC type objects on this platform")
+
+    # Delete the enum type. Returning an instance from Python should fail
+    # rather than accessing a deleted object.
+    pytest.delattr_and_ensure_destroyed((m, "altitude"))
+    with pytest.raises(TypeError, match="Unable to convert function return"):
+        m.get_altitude()
+    with pytest.raises(TypeError, match="incompatible function arguments"):
+        m.is_high_altitude("oops")
+
+    # Recreate the enum type; should not have any duplicate-binding error
+    m.bind_altitude(m)
+    test_altitude_enum()
+    test_altitude_binding()
+
+    # Remove the pybind11 capsule without removing the type; enum is still
+    # usable but can't be passed to/from bound functions
+    del m.altitude.__pybind11_native_enum__
+    pytest.gc_collect()
+    test_altitude_enum()  # enum itself still works
+
+    with pytest.raises(TypeError, match="Unable to convert function return"):
+        m.get_altitude()
+    with pytest.raises(TypeError, match="incompatible function arguments"):
+        m.is_high_altitude(m.altitude.high)
+
+    del m.altitude