pyo3_pytests/
enums.rs

1use pyo3::{
2    pyclass, pyfunction, pymodule,
3    types::{PyModule, PyModuleMethods},
4    wrap_pyfunction, Bound, PyResult,
5};
6
7#[pymodule(gil_used = false)]
8pub fn enums(m: &Bound<'_, PyModule>) -> PyResult<()> {
9    m.add_class::<SimpleEnum>()?;
10    m.add_class::<ComplexEnum>()?;
11    m.add_class::<SimpleTupleEnum>()?;
12    m.add_class::<TupleEnum>()?;
13    m.add_class::<MixedComplexEnum>()?;
14    m.add_wrapped(wrap_pyfunction!(do_simple_stuff))?;
15    m.add_wrapped(wrap_pyfunction!(do_complex_stuff))?;
16    m.add_wrapped(wrap_pyfunction!(do_tuple_stuff))?;
17    m.add_wrapped(wrap_pyfunction!(do_mixed_complex_stuff))?;
18    Ok(())
19}
20
21#[pyclass(eq, eq_int)]
22#[derive(PartialEq)]
23pub enum SimpleEnum {
24    Sunday,
25    Monday,
26    Tuesday,
27    Wednesday,
28    Thursday,
29    Friday,
30    Saturday,
31}
32
33#[pyfunction]
34pub fn do_simple_stuff(thing: &SimpleEnum) -> SimpleEnum {
35    match thing {
36        SimpleEnum::Sunday => SimpleEnum::Monday,
37        SimpleEnum::Monday => SimpleEnum::Tuesday,
38        SimpleEnum::Tuesday => SimpleEnum::Wednesday,
39        SimpleEnum::Wednesday => SimpleEnum::Thursday,
40        SimpleEnum::Thursday => SimpleEnum::Friday,
41        SimpleEnum::Friday => SimpleEnum::Saturday,
42        SimpleEnum::Saturday => SimpleEnum::Sunday,
43    }
44}
45
46#[pyclass]
47pub enum ComplexEnum {
48    Int {
49        i: i32,
50    },
51    Float {
52        f: f64,
53    },
54    Str {
55        s: String,
56    },
57    EmptyStruct {},
58    MultiFieldStruct {
59        a: i32,
60        b: f64,
61        c: bool,
62    },
63    #[pyo3(constructor = (a = 42, b = None))]
64    VariantWithDefault {
65        a: i32,
66        b: Option<String>,
67    },
68}
69
70#[pyfunction]
71pub fn do_complex_stuff(thing: &ComplexEnum) -> ComplexEnum {
72    match thing {
73        ComplexEnum::Int { i } => ComplexEnum::Str { s: i.to_string() },
74        ComplexEnum::Float { f } => ComplexEnum::Float { f: f * f },
75        ComplexEnum::Str { s } => ComplexEnum::Int { i: s.len() as i32 },
76        ComplexEnum::EmptyStruct {} => ComplexEnum::EmptyStruct {},
77        ComplexEnum::MultiFieldStruct { a, b, c } => ComplexEnum::MultiFieldStruct {
78            a: *a,
79            b: *b,
80            c: *c,
81        },
82        ComplexEnum::VariantWithDefault { a, b } => ComplexEnum::VariantWithDefault {
83            a: 2 * a,
84            b: b.as_ref().map(|s| s.to_uppercase()),
85        },
86    }
87}
88
89#[pyclass]
90enum SimpleTupleEnum {
91    Int(i32),
92    Str(String),
93}
94
95#[pyclass]
96pub enum TupleEnum {
97    #[pyo3(constructor = (_0 = 1, _1 = 1.0, _2 = true))]
98    FullWithDefault(i32, f64, bool),
99    Full(i32, f64, bool),
100    EmptyTuple(),
101}
102
103#[pyfunction]
104pub fn do_tuple_stuff(thing: &TupleEnum) -> TupleEnum {
105    match thing {
106        TupleEnum::FullWithDefault(a, b, c) => TupleEnum::FullWithDefault(*a, *b, *c),
107        TupleEnum::Full(a, b, c) => TupleEnum::Full(*a, *b, *c),
108        TupleEnum::EmptyTuple() => TupleEnum::EmptyTuple(),
109    }
110}
111
112#[pyclass]
113pub enum MixedComplexEnum {
114    Nothing {},
115    Empty(),
116}
117
118#[pyfunction]
119pub fn do_mixed_complex_stuff(thing: &MixedComplexEnum) -> MixedComplexEnum {
120    match thing {
121        MixedComplexEnum::Nothing {} => MixedComplexEnum::Empty(),
122        MixedComplexEnum::Empty() => MixedComplexEnum::Nothing {},
123    }
124}
pyo3_pytests/enums.rs

pyo3_pytests/
enums.rs
