pyo3 列挙型

Rustのpyo3で列挙型を定義する方法のメモ。

概要

Pythonの列挙型に相当するクラスをpyo3 0.27で実現するには、Rustのenumを定義し、pyclass属性を付ける。

use pyo3::prelude::pyclass;

厳密には、Pythonの列挙型はEnumを継承したクラスとして定義するらしいが、PyO3ではEnumを継承する形にはならない。
（Pythonでisinstance(列挙型の値, Enum)を実行するとFalseになる）

例

src/my_enum.rs：

use pyo3::prelude::*;

#[pyclass]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[allow(clippy::upper_case_acronyms)]
pub enum Color {
    RED,
    GREEN,
    BLUE,
}

Rustのコーディング規約では列挙子はCamelCaseだが、PythonではUPPER_CASE（定数扱い）らしいので、大文字で定義しておく方がよさそう。

#[pymethods]
impl Color {
    pub fn __eq__(&self, other: &Color) -> bool {
        self == other
    }

    pub fn __hash__(&self) -> isize {
        *self as isize
    }
}

ついでに、__eq__メソッドを定義しておくといいらしい。
（ColorをRustの別の構造体のフィールドに保持してPythonから取り出すような場合、同じ列挙子でも別インスタンスになるらしい。そうなると、デフォルトでは「==」で比較したときにFalseになってしまう。__eq__メソッドを実装しておくと、「==」の実行時にそれが呼ばれる）

さらについでに、__hash__メソッドを定義しておくと、辞書（dict）のキーとして使えるようになるらしい。

src/lib.rs：

use pyo3::prelude::*;

mod my_enum;

#[pymodule]
mod example_pyo3 {
    use pyo3::prelude::*;

    #[pymodule_export]
    use crate::my_enum::Color;
}

呼び出すPython側の例

from example_pyo3 import Color

    value = Color.RED
    print(value)                # Color.RED
    print(value == Color.RED)   # True
    print(value == Color.BLUE)  # False