Subscriptable enums

[markkimsal]

I have some Python code that uses enums with dynamic subscript string keys. I can't seem to figure out how to define then in Rust and export them to Py. When I try to define __getitem__ it doesn't seem to work because it's passing a Type or Class and not an instance of the object (because it's an enum, not an object).

        # a is a parsed json dictionary
        acct = Account(
            a.get("id", 0),
            AccountType[a.get("type")],
            OwnerType[a.get("owner_type").upper()],
            a.get("balance", 0),

#[derive(Debug)]
#[pyclass]
pub enum AccountType {
    INVESTMENT,
    RETIREMENT,
    IRA,
}

#[pymethods]
impl AccountType {
    fn __getitem__(&self, item: &PyAny) -> PyResult<Self> {
        let item_string = item.extract::<String>().unwrap();
        match item_string.as_str() {
            "INVESTMENT" => Ok(Self::INVESTMENT),
            _ => Err(PyKeyError::new_err("key not found in enum"))
        }
    }
}

When I try to use it form python, it says TypeError: 'type' object is not subscriptable

If I try to call getItem directly it says that ___getitem___ requires a 'builtins.AccountType' object, but received a 'str'

I've tried adding # [staticmethod] to getitem as well.

AccountType['RETIREMENT']
Traceback (most recent call last):
  File "<string>", line 1, in <module>
TypeError: 'type' object is not subscriptable

AccountType.__getitem__("sdf")
Traceback (most recent call last):
  File "<string>", line 1, in <module>
TypeError: descriptor '__getitem__' requires a 'builtins.AccountType' object but received a 'str'

Anyone have an idea how to make this somewhat doable? Like using a dictionary instead of an enum maybe? Or write matching parsing/matching functions in both languages?

[markkimsal]

I think I figured it out with __class_getitem__

    #[staticmethod]
    fn __class_getitem__(item: &str) -> PyResult<Self> {
        match item {
            "INVESTMENT" => Ok(Self::INVESTMENT),
            "RETIREMENT" => Ok(Self::RETIREMENT),
            "IRA"        => Ok(Self::IRA),
            _            => Err(PyKeyError::new_err("key not found in enum")),
        }
    }

Although, this function is not mentioned in #2887 (comment)

[markkimsal]

I guess the __class_getitem__ solution is for the current class based, not-real-enum solution and the previously mentioned issue is about implementing more "real" enums. Is that a correct assumption?

It would be nice to see some validation that this is the right way to go given the current # [pyclass] implementation

[davidhewitt]

Yes this looks like the right approach. Are the original Python enums subclasses of enum.Enum? If so, I guess this is indeed something to add to #2887