pyo3_pytests/
enums.rs

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

pyo3_pytests/
enums.rs
