feat(PyTreeKind): use pybind11::native_enum to create enum class PyTreeKind

Author: XuehaiPan

include/optree/pymacros.h

@@ -21,12 +21,16 @@ limitations under the License.
 
 #include <pybind11/pybind11.h>
 
-namespace py = pybind11;
-
-#if PY_VERSION_HEX < 0x03090000  // Python 3.9
+#if !(defined(PY_VERSION_HEX) && PY_VERSION_HEX >= 0x03090000)  // Python 3.9
 #error "Python 3.9 or newer is required."
 #endif
 
+#if !(defined(PYBIND11_VERSION_HEX) && PYBIND11_VERSION_HEX >= 0x020C00F0)  // pybind11 2.12.0
+#error "pybind11 2.12.0 or newer is required."
+#endif
+
+namespace py = pybind11;
+
 #ifndef Py_ALWAYS_INLINE
 #define Py_ALWAYS_INLINE
 #endif

src/optree.cpp

@@ -25,6 +25,16 @@ limitations under the License.
 #include <pybind11/pybind11.h>
 #include <pybind11/stl.h>
 
+#if defined(PYBIND11_HAS_NATIVE_ENUM) ||                                                           \
+    (defined(PYBIND11_INTERNALS_VERSION) && PYBIND11_INTERNALS_VERSION >= 8)
+#ifndef PYBIND11_HAS_NATIVE_ENUM
+#define PYBIND11_HAS_NATIVE_ENUM
+#endif
+#include <pybind11/native_enum.h>
+#else
+#undef PYBIND11_HAS_NATIVE_ENUM
+#endif
+
 namespace optree {
 
 py::module_ GetCxxModule(const std::optional<py::module_>& module) {
@@ -175,6 +185,22 @@ void BuildModule(py::module_& mod) {  // NOLINT[runtime/references]
 #define def_method_pos_only(...) def(__VA_ARGS__)
 #endif
 
+#ifdef PYBIND11_HAS_NATIVE_ENUM
+    py::native_enum<PyTreeKind>(mod, "PyTreeKind", "enum.IntEnum", "The kind of a pytree node.")
+        .value("CUSTOM", PyTreeKind::Custom, "A custom type.")
+        .value("LEAF", PyTreeKind::Leaf, "A opaque leaf node.")
+        .value("NONE", PyTreeKind::None, "None.")
+        .value("TUPLE", PyTreeKind::Tuple, "A tuple.")
+        .value("LIST", PyTreeKind::List, "A list.")
+        .value("DICT", PyTreeKind::Dict, "A dict.")
+        .value("NAMEDTUPLE", PyTreeKind::NamedTuple, "A collections.namedtuple.")
+        .value("ORDEREDDICT", PyTreeKind::OrderedDict, "A collections.OrderedDict.")
+        .value("DEFAULTDICT", PyTreeKind::DefaultDict, "A collections.defaultdict.")
+        .value("DEQUE", PyTreeKind::Deque, "A collections.deque.")
+        .value("STRUCTSEQUENCE", PyTreeKind::StructSequence, "A PyStructSequence.")
+        .finalize();
+    auto PyTreeKindTypeObject = py::getattr(mod, "PyTreeKind");
+#else
     auto PyTreeKindTypeObject =
         py::enum_<PyTreeKind>(mod, "PyTreeKind", "The kind of a pytree node.", py::module_local())
             .value("CUSTOM", PyTreeKind::Custom, "A custom type.")
@@ -188,6 +214,7 @@ void BuildModule(py::module_& mod) {  // NOLINT[runtime/references]
             .value("DEFAULTDICT", PyTreeKind::DefaultDict, "A collections.defaultdict.")
             .value("DEQUE", PyTreeKind::Deque, "A collections.deque.")
             .value("STRUCTSEQUENCE", PyTreeKind::StructSequence, "A PyStructSequence.");
+#endif
     auto* const PyTreeKind_Type = reinterpret_cast<PyTypeObject*>(PyTreeKindTypeObject.ptr());
     PyTreeKind_Type->tp_name = "optree.PyTreeKind";
     py::setattr(PyTreeKindTypeObject.ptr(), Py_Get_ID(__module__), Py_Get_ID(optree));
@@ -412,25 +439,35 @@ void BuildModule(py::module_& mod) {  // NOLINT[runtime/references]
 
 #undef def_method_pos_only
 
+// Make the PyTreeSpec and PyTreeIter types immutable.
 #ifdef Py_TPFLAGS_IMMUTABLETYPE
-    PyTreeKind_Type->tp_flags |= Py_TPFLAGS_IMMUTABLETYPE;
     PyTreeSpec_Type->tp_flags |= Py_TPFLAGS_IMMUTABLETYPE;
-    PyTreeIter_Type->tp_flags |= Py_TPFLAGS_IMMUTABLETYPE;
-    PyTreeKind_Type->tp_flags &= ~Py_TPFLAGS_READY;
     PyTreeSpec_Type->tp_flags &= ~Py_TPFLAGS_READY;
+    PyTreeIter_Type->tp_flags |= Py_TPFLAGS_IMMUTABLETYPE;
     PyTreeIter_Type->tp_flags &= ~Py_TPFLAGS_READY;
 #endif
 
-    if (PyType_Ready(PyTreeKind_Type) < 0) [[unlikely]] {
-        INTERNAL_ERROR("`PyType_Ready(&PyTreeKind_Type)` failed.");
-    }
     if (PyType_Ready(PyTreeSpec_Type) < 0) [[unlikely]] {
         INTERNAL_ERROR("`PyType_Ready(&PyTreeSpec_Type)` failed.");
     }
     if (PyType_Ready(PyTreeIter_Type) < 0) [[unlikely]] {
         INTERNAL_ERROR("`PyType_Ready(&PyTreeIter_Type)` failed.");
     }
 
+#ifdef Py_TPFLAGS_IMMUTABLETYPE
+    PyTreeKind_Type->tp_flags |= Py_TPFLAGS_IMMUTABLETYPE;
+
+    // Only run `PyType_Ready` for C++ types.
+    // Re-run `PyType_Ready` for native Python enums will cause unexpected behavior.
+#ifndef PYBIND11_HAS_NATIVE_ENUM
+    PyTreeKind_Type->tp_flags &= ~Py_TPFLAGS_READY;
+
+    if (PyType_Ready(PyTreeKind_Type) < 0) [[unlikely]] {
+        INTERNAL_ERROR("`PyType_Ready(&PyTreeKind_Type)` failed.");
+    }
+#endif
+#endif
+
     py::getattr(py::module_::import("atexit"),
                 "register")(py::cpp_function(&PyTreeTypeRegistry::Clear));
 }