1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81

use pyo3::{
    pyclass, pyfunction, pymodule,
    types::{PyModule, PyModuleMethods},
    wrap_pyfunction_bound, Bound, PyResult,
};

#[pymodule]
pub fn enums(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<SimpleEnum>()?;
    m.add_class::<ComplexEnum>()?;
    m.add_wrapped(wrap_pyfunction_bound!(do_simple_stuff))?;
    m.add_wrapped(wrap_pyfunction_bound!(do_complex_stuff))?;
    Ok(())
}

#[pyclass]
pub enum SimpleEnum {
    Sunday,
    Monday,
    Tuesday,
    Wednesday,
    Thursday,
    Friday,
    Saturday,
}

#[pyfunction]
pub fn do_simple_stuff(thing: &SimpleEnum) -> SimpleEnum {
    match thing {
        SimpleEnum::Sunday => SimpleEnum::Monday,
        SimpleEnum::Monday => SimpleEnum::Tuesday,
        SimpleEnum::Tuesday => SimpleEnum::Wednesday,
        SimpleEnum::Wednesday => SimpleEnum::Thursday,
        SimpleEnum::Thursday => SimpleEnum::Friday,
        SimpleEnum::Friday => SimpleEnum::Saturday,
        SimpleEnum::Saturday => SimpleEnum::Sunday,
    }
}

#[pyclass]
pub enum ComplexEnum {
    Int {
        i: i32,
    },
    Float {
        f: f64,
    },
    Str {
        s: String,
    },
    EmptyStruct {},
    MultiFieldStruct {
        a: i32,
        b: f64,
        c: bool,
    },
    #[pyo3(constructor = (a = 42, b = None))]
    VariantWithDefault {
        a: i32,
        b: Option<String>,
    },
}

#[pyfunction]
pub fn do_complex_stuff(thing: &ComplexEnum) -> ComplexEnum {
    match thing {
        ComplexEnum::Int { i } => ComplexEnum::Str { s: i.to_string() },
        ComplexEnum::Float { f } => ComplexEnum::Float { f: f * f },
        ComplexEnum::Str { s } => ComplexEnum::Int { i: s.len() as i32 },
        ComplexEnum::EmptyStruct {} => ComplexEnum::EmptyStruct {},
        ComplexEnum::MultiFieldStruct { a, b, c } => ComplexEnum::MultiFieldStruct {
            a: *a,
            b: *b,
            c: *c,
        },
        ComplexEnum::VariantWithDefault { a, b } => ComplexEnum::VariantWithDefault {
            a: 2 * a,
            b: b.as_ref().map(|s| s.to_uppercase()),
        },
    }
}