Add native_enum: expose Rust enums as Python enum.Enum subclasses

newsfragments/6020.added.md

@@ -0,0 +1 @@
+Added `#[py_native_enum]` / `#[derive(NativeEnum)]` to expose fieldless Rust enums as Python `enum.Enum` subclasses, supporting `Enum`, `IntEnum`, `StrEnum`, `Flag`, and `IntFlag` bases.

pyo3-benches/Cargo.toml

@@ -86,6 +86,10 @@ harness = false
 name = "bench_intern"
 harness = false
 
+[[bench]]
+name = "bench_native_enum"
+harness = false
+
 [[bench]]
 name = "bench_extract"
 harness = false

pyo3-benches/benches/bench_native_enum.rs

@@ -0,0 +1,74 @@
+use std::hint::black_box;
+
+use codspeed_criterion_compat::{criterion_group, criterion_main, Bencher, Criterion};
+
+use pyo3::native_enum::NativeEnum;
+use pyo3::prelude::*;
+use pyo3::py_native_enum;
+
+/// A simple enum using the default `enum.Enum` base.
+#[py_native_enum]
+enum Color {
+    Red,
+    Green,
+    Blue,
+}
+
+/// An integer enum using `enum.IntEnum`.
+#[py_native_enum(base = "IntEnum")]
+enum Status {
+    Active,
+    Inactive,
+    Pending,
+}
+
+// Measures the PyOnceLock cache-hit path: get + clone_ref + into_bound.
+fn bench_py_enum_class(b: &mut Bencher<'_>) {
+    Python::attach(|py| {
+        b.iter(|| black_box(Color::py_enum_class(py).unwrap()));
+    });
+}
+
+fn bench_int_enum_class(b: &mut Bencher<'_>) {
+    Python::attach(|py| {
+        b.iter(|| black_box(Status::py_enum_class(py).unwrap()));
+    });
+}
+
+fn bench_to_py_member(b: &mut Bencher<'_>) {
+    Python::attach(|py| {
+        b.iter_with_large_drop(|| Color::Green.to_py_member(py).unwrap());
+    });
+}
+
+fn bench_int_enum_to_py_member(b: &mut Bencher<'_>) {
+    Python::attach(|py| {
+        b.iter_with_large_drop(|| Status::Active.to_py_member(py).unwrap());
+    });
+}
+
+fn bench_from_py_member(b: &mut Bencher<'_>) {
+    Python::attach(|py| {
+        let obj = Color::Blue.to_py_member(py).unwrap();
+        b.iter(|| Color::from_py_member(black_box(&obj)).unwrap());
+    });
+}
+
+fn bench_int_enum_from_py_member(b: &mut Bencher<'_>) {
+    Python::attach(|py| {
+        let obj = Status::Pending.to_py_member(py).unwrap();
+        b.iter(|| Status::from_py_member(black_box(&obj)).unwrap());
+    });
+}
+
+fn criterion_benchmark(c: &mut Criterion) {
+    c.bench_function("native_enum_py_enum_class", bench_py_enum_class);
+    c.bench_function("native_enum_int_enum_class", bench_int_enum_class);
+    c.bench_function("native_enum_to_py_member", bench_to_py_member);
+    c.bench_function("native_enum_int_enum_to_py_member", bench_int_enum_to_py_member);
+    c.bench_function("native_enum_from_py_member", bench_from_py_member);
+    c.bench_function("native_enum_int_enum_from_py_member", bench_int_enum_from_py_member);
+}
+
+criterion_group!(benches, criterion_benchmark);
+criterion_main!(benches);

pyo3-macros-backend/src/attributes.rs

@@ -56,6 +56,8 @@ pub mod kw {
     syn::custom_keyword!(category);
     syn::custom_keyword!(from_py_object);
     syn::custom_keyword!(skip_from_py_object);
+    syn::custom_keyword!(base);
+    syn::custom_keyword!(value);
 }
 
 fn take_int(read: &mut &str, tracker: &mut usize) -> String {

pyo3-macros-backend/src/lib.rs

@@ -18,6 +18,7 @@ mod introspection;
 mod konst;
 mod method;
 mod module;
+mod native_enum;
 mod params;
 #[cfg(feature = "experimental-inspect")]
 mod py_expr;
@@ -30,6 +31,7 @@ mod quotes;
 pub use frompyobject::build_derive_from_pyobject;
 pub use intopyobject::build_derive_into_pyobject;
 pub use module::{pymodule_function_impl, pymodule_module_impl, PyModuleOptions};
+pub use native_enum::{build_derive_native_enum, native_enum_impl, PyNativeEnumArgs};
 pub use pyclass::{build_py_class, build_py_enum, PyClassArgs};
 pub use pyfunction::{build_py_function, PyFunctionOptions};
 pub use pyimpl::{build_py_methods, PyClassMethodsType};