revived enum::def, def_prop_ro, and def_prop_rw

wjakob · May 9, 2024 · 2889d77 · 2889d77
1 parent 12c3c79
commit 2889d77
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 1 deletion.
diff --git a/include/nanobind/nb_class.h b/include/nanobind/nb_class.h
@@ -721,6 +721,40 @@ template <typename T> class enum_ : public object {
     }
 
     NB_INLINE enum_ &export_values() { detail::enum_export(m_ptr); return *this; }
+
+
+    template <typename Func, typename... Extra>
+    NB_INLINE enum_ &def(const char *name_, Func &&f, const Extra &... extra) {
+        cpp_function_def<T>((detail::forward_t<Func>) f, scope(*this),
+                            name(name_), is_method(), extra...);
+        return *this;
+    }
+
+    template <typename Getter, typename Setter, typename... Extra>
+    NB_INLINE enum_ &def_prop_rw(const char *name_, Getter &&getter,
+                                 Setter &&setter, const Extra &...extra) {
+        object get_p, set_p;
+
+        if constexpr (!std::is_same_v<Getter, std::nullptr_t>)
+            get_p = cpp_function<T>((detail::forward_t<Getter>) getter,
+                                    is_method(), is_getter(),
+                                    rv_policy::reference_internal,
+                                    detail::filter_getter(extra)...);
+
+        if constexpr (!std::is_same_v<Setter, std::nullptr_t>)
+            set_p = cpp_function<T>((detail::forward_t<Setter>) setter,
+                                    is_method(), detail::filter_setter(extra)...);
+
+        detail::property_install(m_ptr, name_, get_p.ptr(), set_p.ptr());
+        return *this;
+    }
+
+
+    template <typename Getter, typename... Extra>
+    NB_INLINE enum_ &def_prop_ro(const char *name_, Getter &&getter,
+                                 const Extra &...extra) {
+        return def_prop_rw(name_, getter, nullptr, extra...);
+    }
 };
 
 template <typename Source, typename Target> void implicitly_convertible() {

diff --git a/tests/test_enum.cpp b/tests/test_enum.cpp
@@ -19,11 +19,14 @@ NB_MODULE(test_enum_ext, m) {
         .value("B", SEnum::B)
         .value("C", SEnum::C);
 
-    nb::enum_<ClassicEnum>(m, "ClassicEnum")
+    auto ce = nb::enum_<ClassicEnum>(m, "ClassicEnum")
         .value("Item1", ClassicEnum::Item1)
         .value("Item2", ClassicEnum::Item2)
         .export_values();
 
+    ce.def("get_value", [](ClassicEnum &x) { return (int) x; })
+      .def_prop_ro("my_value", [](ClassicEnum &x) { return (int) x; });
+
     m.def("from_enum", [](Enum value) { return (uint32_t) value; }, nb::arg().noconvert());
     m.def("to_enum", [](uint32_t value) { return (Enum) value; });
     m.def("from_enum", [](SEnum value) { return (int32_t) value; }, nb::arg().noconvert());

diff --git a/tests/test_enum.py b/tests/test_enum.py
@@ -135,3 +135,7 @@ def test08_enum_comparisons():
     assert t.Enum.B != t.SEnum.B and t.SEnum.B != t.Enum.B
     assert not (t.Enum.B == t.SEnum.B) and not (t.SEnum.B == t.Enum.B)
     assert t.Enum.B != t.SEnum.C and t.SEnum.C != t.Enum.B
+
+def test09_enum_methods():
+    assert t.Item1.my_value == 0 and t.Item2.my_value == 1
+    assert t.Item1.get_value() == 0 and t.Item2.get_value() == 1