Skip to content

Commit 96d099f

Browse files
committed
feat: support pyclass on complex enums
1 parent 48e74b7 commit 96d099f

File tree

14 files changed

+991
-67
lines changed

14 files changed

+991
-67
lines changed

guide/src/class.md

Lines changed: 97 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22

33
PyO3 exposes a group of attributes powered by Rust's proc macro system for defining Python classes as Rust structs.
44

5-
The main attribute is `#[pyclass]`, which is placed upon a Rust `struct` or a fieldless `enum` (a.k.a. C-like enum) to generate a Python type for it. They will usually also have *one* `#[pymethods]`-annotated `impl` block for the struct, which is used to define Python methods and constants for the generated Python type. (If the [`multiple-pymethods`] feature is enabled, each `#[pyclass]` is allowed to have multiple `#[pymethods]` blocks.) `#[pymethods]` may also have implementations for Python magic methods such as `__str__`.
5+
The main attribute is `#[pyclass]`, which is placed upon a Rust `struct` or `enum` to generate a Python type for it. They will usually also have *one* `#[pymethods]`-annotated `impl` block for the struct, which is used to define Python methods and constants for the generated Python type. (If the [`multiple-pymethods`] feature is enabled, each `#[pyclass]` is allowed to have multiple `#[pymethods]` blocks.) `#[pymethods]` may also have implementations for Python magic methods such as `__str__`.
66

77
This chapter will discuss the functionality and configuration these attributes offer. Below is a list of links to the relevant section of this chapter for each:
88

@@ -21,21 +21,29 @@ This chapter will discuss the functionality and configuration these attributes o
2121

2222
## Defining a new class
2323

24-
To define a custom Python class, add the `#[pyclass]` attribute to a Rust struct or a fieldless enum.
24+
To define a custom Python class, add the `#[pyclass]` attribute to a Rust struct or enum.
2525
```rust
2626
# #![allow(dead_code)]
2727
use pyo3::prelude::*;
2828

2929
#[pyclass]
30-
struct Integer {
30+
struct MyClass {
3131
inner: i32,
3232
}
3333

3434
// A "tuple" struct
3535
#[pyclass]
3636
struct Number(i32);
3737

38-
// PyO3 supports custom discriminants in enums
38+
// PyO3 supports unit-only enums (which contain only unit variants)
39+
// These simple enums behave similarly to Python's enumerations (enum.Enum)
40+
#[pyclass]
41+
enum MyEnum {
42+
Variant,
43+
OtherVariant = 30, // PyO3 supports custom discriminants.
44+
}
45+
46+
// PyO3 supports custom discriminants in unit-only enums
3947
#[pyclass]
4048
enum HttpResponse {
4149
Ok = 200,
@@ -44,14 +52,19 @@ enum HttpResponse {
4452
// ...
4553
}
4654

55+
// PyO3 also supports enums with non-unit variants
56+
// These complex enums have sligtly different behavior from the simple enums above
57+
// They are meant to work with instance checks and match statement patterns
4758
#[pyclass]
48-
enum MyEnum {
49-
Variant,
50-
OtherVariant = 30, // PyO3 supports custom discriminants.
59+
enum Shape {
60+
Circle { radius: f64 },
61+
Rectangle { width: f64, height: f64 },
62+
RegularPolygon { side_count: u32, radius: f64 },
63+
Nothing { },
5164
}
5265
```
5366

54-
The above example generates implementations for [`PyTypeInfo`] and [`PyClass`] for `MyClass` and `MyEnum`. To see these generated implementations, refer to the [implementation details](#implementation-details) at the end of this chapter.
67+
The above example generates implementations for [`PyTypeInfo`] and [`PyClass`] for `MyClass`, `Number`, `MyEnum`, `HttpResponse`, and `Shape`. To see these generated implementations, refer to the [implementation details](#implementation-details) at the end of this chapter.
5568

5669
### Restrictions
5770

@@ -964,7 +977,13 @@ Note that `text_signature` on `#[new]` is not compatible with compilation in
964977

965978
## #[pyclass] enums
966979

967-
Currently PyO3 only supports fieldless enums. PyO3 adds a class attribute for each variant, so you can access them in Python without defining `#[new]`. PyO3 also provides default implementations of `__richcmp__` and `__int__`, so they can be compared using `==`:
980+
Enum support in PyO3 comes in two flavors, depending on what kind of variants the enum has: simple and complex.
981+
982+
### Simple enums
983+
984+
A simple enum (a.k.a. C-like enum) has only unit variants.
985+
986+
PyO3 adds a class attribute for each variant, so you can access them in Python without defining `#[new]`. PyO3 also provides default implementations of `__richcmp__` and `__int__`, so they can be compared using `==`:
968987

969988
```rust
970989
# use pyo3::prelude::*;
@@ -986,7 +1005,7 @@ Python::with_gil(|py| {
9861005
})
9871006
```
9881007

989-
You can also convert your enums into `int`:
1008+
You can also convert your simple enums into `int`:
9901009

9911010
```rust
9921011
# use pyo3::prelude::*;
@@ -1094,6 +1113,74 @@ enum BadSubclass {
10941113

10951114
`#[pyclass]` enums are currently not interoperable with `IntEnum` in Python.
10961115

1116+
### Complex enums
1117+
1118+
An enum is complex if it has any non-unit (struct or tuple) variants.
1119+
1120+
Currently PyO3 supports only struct variants in a complex enum. Support for unit and tuple variants is planned.
1121+
1122+
PyO3 adds a class attribute for each variant, which may be used to construct values and in match patterns. PyO3 also provides getter methods for all fields of each variant.
1123+
1124+
```rust
1125+
# use pyo3::prelude::*;
1126+
#[pyclass]
1127+
enum Shape {
1128+
Circle { radius: f64 },
1129+
Rectangle { width: f64, height: f64 },
1130+
RegularPolygon { side_count: u32, radius: f64 },
1131+
Nothing { },
1132+
}
1133+
1134+
Python::with_gil(|py| {
1135+
let circle = Shape::Circle { radius: 10.0 }.into_py(py);
1136+
let square = Shape::RegularPolygon { side_count: 4, radius: 10.0 }.into_py(py);
1137+
let cls = py.get_type::<Shape>();
1138+
pyo3::py_run!(py, circle square cls, r#"
1139+
assert isinstance(circle, cls)
1140+
assert isinstance(circle, cls.Circle)
1141+
assert circle.radius == 10.0
1142+
1143+
assert isinstance(square, cls)
1144+
assert isinstance(square, cls.RegularPolygon)
1145+
assert square.side_count == 4
1146+
assert square.radius == 10.0
1147+
1148+
def count_vertices(cls, shape):
1149+
match shape:
1150+
case cls.Circle(radius=r):
1151+
return 0
1152+
case cls.Rectangle(width=w, height=h):
1153+
return 4
1154+
case cls.RegularPolygon(side_count=n, radius=r):
1155+
return n
1156+
case cls.Nothing():
1157+
return 0
1158+
1159+
assert count_vertices(cls, circle) == 0
1160+
assert count_vertices(cls, square) == 4
1161+
"#)
1162+
})
1163+
```
1164+
1165+
WARNING: `Py::new` and `.into_py` are currently inconsistent. Note how the constructed value is _not_ an instance of the specific variant. For this reason, constructing values is only recommended using `.into_py`.
1166+
1167+
```rust
1168+
# use pyo3::prelude::*;
1169+
#[pyclass]
1170+
enum MyEnum {
1171+
Variant { i: i32 },
1172+
}
1173+
1174+
Python::with_gil(|py| {
1175+
let x = Py::new(py, MyEnum::Variant { i: 42 }).unwrap();
1176+
let cls = py.get_type::<MyEnum>();
1177+
pyo3::py_run!(py, x cls, r#"
1178+
assert isinstance(x, cls)
1179+
assert not isinstance(x, cls.Variant)
1180+
"#)
1181+
})
1182+
```
1183+
10971184
## Implementation details
10981185

10991186
The `#[pyclass]` macros rely on a lot of conditional code generation: each `#[pyclass]` can optionally have a `#[pymethods]` block.

newsfragments/3582.added.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Support `#[pyclass]` on enums that have non-unit variants.

0 commit comments

Comments
 (0)