feat(treespec): set enum type immutable

metaopt · Oct 28, 2023 · fd156b6 · fd156b6
1 parent ae03919
commit fd156b6
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/src/optree.cpp b/src/optree.cpp
@@ -251,11 +251,17 @@ void BuildModule(py::module& mod) {  // NOLINT[runtime/references]
 #ifdef Py_TPFLAGS_IMMUTABLETYPE
     reinterpret_cast<PyTypeObject*>(PyTreeSpecTypeObject.ptr())->tp_flags |=
         Py_TPFLAGS_IMMUTABLETYPE;
+    reinterpret_cast<PyTypeObject*>(PyTreeKindTypeObject.ptr())->tp_flags |=
+        Py_TPFLAGS_IMMUTABLETYPE;
 #endif
     if (PyType_Ready(reinterpret_cast<PyTypeObject*>(PyTreeSpecTypeObject.ptr())) < 0)
         [[unlikely]] {
         INTERNAL_ERROR("`PyType_Ready(&PyTreeSpec_Type)` failed.");
     }
+    if (PyType_Ready(reinterpret_cast<PyTypeObject*>(PyTreeKindTypeObject.ptr())) < 0)
+        [[unlikely]] {
+        INTERNAL_ERROR("`PyType_Ready(&PyTreeKind_Type)` failed.");
+    }
 }
 
 }  // namespace optree