PyO3 - How to return enums to python module?

Question

I'm trying to build a Python package from Rust using PyO3. Right now I'm stuck trying to return enums Rust type to Python.

I have a simple enum like so:

pub enum Lang {
    Deu,
    Eng,
    Fra
}

And in lib.rs

#[pyfunction]
fn detect_language(text: &str) -> PyResult<????> {
   // Do some stuff ....
   res:Lang = Do_some_stuff(text)
   Ok(res)
}

#[pymodule]
fn pymylib(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(detect_language, m)?)?;
    Ok(())
}

In Python code

from pymylib import detect_language

res=detect_language('Ceci est un test')
print(res) # Lang:Fra ???

Abdul Niyas P M · Accepted Answer · 2023-05-24T10:45:10.860

One approach would to use #[pyclass] attribute to make the Python class from Rust Enum. Also make sure to export this class from Rust code, so that you can do the comparison on python layer. ie,

use pyo3::prelude::*;

#[pyclass]
pub enum Lang {
    Deu,
    Eng,
    Fra
}

#[pyfunction]
fn detect_language(text: &str) -> PyResult<Lang> {
    // Write your actual code here
    // But for testing purpose let's return `Deu` varient.
    Ok(Lang::Deu)
}

/// A Python module implemented in Rust.
#[pymodule]
fn pymylib(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(detect_language, m)?)?;
    m.add_class::<Lang>()?;
    Ok(())
}

Now in Python interpreter you can do,

>>> from pymylib import Lang
>>> import pymylib
>>> 
>>> pymylib.detect_language("1")
Lang.Deu
>>> pymylib.detect_language("1") == Lang.Deu
True

See also

Support exporting Rust enums to Python

Thank you for your answer. But if I have a variable of type `enum` defined in another crate, how can use #[pyclass] attribute to make the Python class from this variable? eg: ``` use cratelang; #[pyclass] type Lang = cratelang::Lang; // Got ERROR ``` — LeMoussel, May 24 '23 at 10:38

score 1 · Answer 2 · answered May 24 '23 at 11:08

In combination with the fact that python's naming guidelines for enum variants are different to rust's, I generally turn the rust enum into a string in the pyo3 interface, and then wrap it in a StrEnum on the python side. This doesn't require too much extra code:

// interface.rs

#[pyfunction]
fn detect_language(text: &str) -> PyResult<&'static str> {
    Ok(match actual_detect_language(text) {
        Lang::Deu => "deu",
        Lang::Eng => "eng",
        Lang::Fra => "fra",
    })
}

# interface.py
from enum import StrEnum

import pymylib

class Languge(StrEnum):
    DEU = "deu"
    ENG = "eng"
    FRA = "fra"

def detect_language(text: str) -> Language:
    return Language(pymylib.detect_language(text))

You can do the same thing in reverse if you need the enum back on the rust side: StrEnum inherits from str and so can be passed straight into any pyo3 function which takes a str. You can just panic on the "other" variant in the match on the rust side because using the StrEnum guarantees the variant exists.

Indeed, this is a forward-facing answer. There is a backport package available for python 3.8+, and the implementation is pretty simple otherwise. — Chris L. Barnes, May 24 '23 at 12:12

PyO3 - How to return enums to python module?

2 Answers2