Comparing enums by identity

[ricohageman]

In python the convention is to compare enums by identity (Enum.A is Enum.A and Enum.A is not Enum.B) instead of equality (Enum.A == Enum.A and Enum.A != Enum.B). However, a straight forward implementation of enums in rust exposed to python using pyo3 does not support this. The documentation also only mentions comparison by equality.

For the particular use-case I'm working with, it is used like the following:

#[pyclass]
pub enum PyEnum {
   A,
   B,
   C,
}

#[pyclass(get_all)]
pub struct PyEnumHoldingStruct {
   py_enum: PyEnum,
   ...
}

#[pymethods]
impl PyEnumHoldingStruct {
   #[new]
   pub fn new(py_enum: PyEnum, ...) -> Self {
      Self {
         py_enum,
         ...
      }
}

#[pymodule]
fn py_module(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_class::<PyEnum>()?;
    m.add_class::<PyEnumHoldingStruct>()?;
}

import py_module

py_a = py_module.PyEnumHoldingStruct(
   py_module.PyEnum.A,
   ...
)

assert py_a.py_enum is py_module.PyEnum.A
assert py_a.py_enum is not py_module.PyEnum.B
assert py_a.py_enum is not py_module.PyEnum.C

In #2384 there is a description on how to use GILOnceCell and #[classattr] to return pointers to the same instance of an enum such that one can compare the enums by identity. Ideally there would be a way to achieve this result without that amount of boilerplate code per enum variant. Would it be possible to automatically generate this means of some setting in #[pyclass]?

I would be able to implement this if there is some guidance and a clear idea on how to support it.

[davidhewitt]

Hello, I agree with you that this is important; I've been aware of this shortcoming in PyO3's enum implementation and while I've wanted to fix it for a while I haven't had any time available to invest in this particular issue.

Your offer to help implement is appreciated, I'd be happy to feedback on the design choices with you and help steer an implementation.

To point you in the right direction, there's a couple of inter-related pieces which would be worth reading and thinking about:

• The macro machinery for #[pyclass] (including on enum items) is in pyo3-macros-backend/src/pyclass.rs. Ultimately if we want to define class attributes for all enum variants, we'll need to add to that code. In particular enum_default_methods is the code which defines the autogenerated attributes which return new instances (i.e. not of identity). Maybe the issue is there, or maybe we should be generating a different IntoPy implementation for C-like enums which returns existing objects.
• Allow #[new] to return existing objects #2384 is likely related, it is a downside of the way we have currently implemented #[new] which prevents us from returning existing objects. If we could return existing objects from #[new], probably boilerplate would already be reduced.

[ricohageman]

or maybe we should be generating a different IntoPy implementation for C-like enums which returns existing objects

I believe this is the way to go, otherwise it's required to use the class attributes everywhere where the enum is created. See the test I've added in 82353cc. Now it's possible to use MyEnum::Variant and not the pyo3 generated class attributes and still be able to compare the enum by identity.