Skip to content

Commit d1d220f

Browse files
committed
Add native_enum: expose Rust enums as Python enum.Enum subclasses
Introduce `#[py_native_enum]` / `#[derive(NativeEnum)]` macros that expose a fieldless Rust enum to Python as a true `enum.Enum` subclass via the functional API, without touching the C-API enum machinery (#991). Runtime (`src/native_enum/`): - `spec.rs` – `NativeEnumSpec`, `NativeEnumBase`, `VariantValue` - `base_cache.rs` – per-interpreter `PyOnceLock` cache for base classes (`enum.Enum`, `IntEnum`, `StrEnum`, `Flag`, `IntFlag`) and `enum.auto` - `construct.rs` – `build_native_enum`: builds a Python enum subclass from a spec without caching - `trait_def.rs` – `NativeEnum` trait with `py_enum_class`, `to_py_member`, `from_py_member` Macro (`pyo3-macros-backend/src/native_enum.rs`): - Parses enum-level attrs (`base`, `rename`, `module`) and variant-level attrs (`rename`, `value`); supports integer discriminants - Caches the Python class in a function-local `static PyOnceLock<Py<PyType>>` in the generated `py_enum_class`, constructing it once per interpreter session - `from_py_member` checks `isinstance(obj, cached_class)` for type safety - Auto-derives `IntoPyObject`, `IntoPyObject for &T`, `FromPyObject` Public API additions to `pyo3-macros/src/lib.rs`: - `#[proc_macro_derive(NativeEnum, attributes(native_enum))]` - `#[proc_macro_attribute] py_native_enum`
1 parent f57bda7 commit d1d220f

12 files changed

Lines changed: 852 additions & 4 deletions

File tree

pyo3-benches/Cargo.toml

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -86,6 +86,10 @@ harness = false
8686
name = "bench_intern"
8787
harness = false
8888

89+
[[bench]]
90+
name = "bench_native_enum"
91+
harness = false
92+
8993
[[bench]]
9094
name = "bench_extract"
9195
harness = false
Lines changed: 74 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,74 @@
1+
use std::hint::black_box;
2+
3+
use codspeed_criterion_compat::{criterion_group, criterion_main, Bencher, Criterion};
4+
5+
use pyo3::native_enum::NativeEnum;
6+
use pyo3::prelude::*;
7+
use pyo3::py_native_enum;
8+
9+
/// A simple enum using the default `enum.Enum` base.
10+
#[py_native_enum]
11+
enum Color {
12+
Red,
13+
Green,
14+
Blue,
15+
}
16+
17+
/// An integer enum using `enum.IntEnum`.
18+
#[py_native_enum(base = "IntEnum")]
19+
enum Status {
20+
Active,
21+
Inactive,
22+
Pending,
23+
}
24+
25+
// Measures the PyOnceLock cache-hit path: get + clone_ref + into_bound.
26+
fn bench_py_enum_class(b: &mut Bencher<'_>) {
27+
Python::attach(|py| {
28+
b.iter(|| Color::py_enum_class(py).unwrap());
29+
});
30+
}
31+
32+
fn bench_int_enum_class(b: &mut Bencher<'_>) {
33+
Python::attach(|py| {
34+
b.iter(|| Status::py_enum_class(py).unwrap());
35+
});
36+
}
37+
38+
fn bench_to_py_member(b: &mut Bencher<'_>) {
39+
Python::attach(|py| {
40+
b.iter_with_large_drop(|| Color::Green.to_py_member(py).unwrap());
41+
});
42+
}
43+
44+
fn bench_int_enum_to_py_member(b: &mut Bencher<'_>) {
45+
Python::attach(|py| {
46+
b.iter_with_large_drop(|| Status::Active.to_py_member(py).unwrap());
47+
});
48+
}
49+
50+
fn bench_from_py_member(b: &mut Bencher<'_>) {
51+
Python::attach(|py| {
52+
let obj = Color::Blue.to_py_member(py).unwrap();
53+
b.iter(|| Color::from_py_member(black_box(&obj)).unwrap());
54+
});
55+
}
56+
57+
fn bench_int_enum_from_py_member(b: &mut Bencher<'_>) {
58+
Python::attach(|py| {
59+
let obj = Status::Pending.to_py_member(py).unwrap();
60+
b.iter(|| Status::from_py_member(black_box(&obj)).unwrap());
61+
});
62+
}
63+
64+
fn criterion_benchmark(c: &mut Criterion) {
65+
c.bench_function("native_enum_py_enum_class", bench_py_enum_class);
66+
c.bench_function("native_enum_int_enum_class", bench_int_enum_class);
67+
c.bench_function("native_enum_to_py_member", bench_to_py_member);
68+
c.bench_function("native_enum_int_enum_to_py_member", bench_int_enum_to_py_member);
69+
c.bench_function("native_enum_from_py_member", bench_from_py_member);
70+
c.bench_function("native_enum_int_enum_from_py_member", bench_int_enum_from_py_member);
71+
}
72+
73+
criterion_group!(benches, criterion_benchmark);
74+
criterion_main!(benches);

pyo3-macros-backend/src/attributes.rs

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,9 @@ pub mod kw {
5656
syn::custom_keyword!(category);
5757
syn::custom_keyword!(from_py_object);
5858
syn::custom_keyword!(skip_from_py_object);
59+
syn::custom_keyword!(base);
60+
syn::custom_keyword!(rename);
61+
syn::custom_keyword!(value);
5962
}
6063

6164
fn take_int(read: &mut &str, tracker: &mut usize) -> String {

pyo3-macros-backend/src/lib.rs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ mod introspection;
1818
mod konst;
1919
mod method;
2020
mod module;
21+
mod native_enum;
2122
mod params;
2223
#[cfg(feature = "experimental-inspect")]
2324
mod py_expr;
@@ -30,6 +31,7 @@ mod quotes;
3031
pub use frompyobject::build_derive_from_pyobject;
3132
pub use intopyobject::build_derive_into_pyobject;
3233
pub use module::{pymodule_function_impl, pymodule_module_impl, PyModuleOptions};
34+
pub use native_enum::{build_derive_native_enum, native_enum_impl, PyNativeEnumArgs};
3335
pub use pyclass::{build_py_class, build_py_enum, PyClassArgs};
3436
pub use pyfunction::{build_py_function, PyFunctionOptions};
3537
pub use pyimpl::{build_py_methods, PyClassMethodsType};

0 commit comments

Comments
 (0)