Enum-like user classes

[orenbenkiki]

Feature request: Allow annotating a class as enum-like for mypy type checking purposes.

I have classes dealing with a value-like (immutable) type.
One of the classes in that family is an enum-like class, whose values are of the new immutable type.
I tried to implement this by deriving from the standard library Enum, but it has pretty strong opinions about the implementation which do not match the way my value type works (it is too different from a plain integer).

No worries, I just implemented my own MyEnumBase base class (and MyEnumMeta for it), loosely inspired by Enum.
It works fine.

The problem is convincing mypy to type check it.
Specifically, if I write:

class EnumExample(MyEnumBase):
    FOO = MyAuto()

Then mypy will not recognize that EnumExample.FOO has the type EnumExample.
I can work around this by the following horrible hack:

class IsEnumBase:
    """Check whether a given class is an enumerated base class."""
    def __eq__(self, class_name: Any) -> bool:
        return class_name == 'enum.Enum' or class_name == 'my.module.MyEnumBase'

mypy.nodes.ENUM_BASECLASS = IsEnumBase()

Which mostly works, but, shudder (is there a better way to achieve the same effect?)

It would be nice if I could annotate MyEnumBase somehow to let mypy it is an enum base class.
For example, if the single mypy.nodes.ENUM_BASECLASS was replaced by a set of class names, so that it were possible to insert additional class names into it (and provide a command line flag to mypy for doing that).

[emmatyping]

Have you tried inheriting from EnumMeta? I believe that was designed to be more customizable.

[ilevkivskyi]

Can this be achieved by descriptors? mypy already supports __get__ and __set__.

[orenbenkiki]

@ethanhs: I did try deriving from EnumMeta, but that was using Python 3.6.0. Perhaps it was modified since but I'm still supporting 3.6.0. Even if I could use it, from looking at the code in mypy.nodes, it seems it explicitly checks whether the base class name is enum.Enum rather than checking whether the base class uses EnumMeta as a metaclass, so I doubt that would solve the type checking problem.

@ilevkivskyi: Not sure what you mean. My problem is in convincing mypy that ExampleEnum.FOO has the actual type it has at run-time (that is, ExampleEnum) rather than the type mypy thinks it has (that is, MyAuto). Since I explicitly initialize the class member, I don't think __set__ and __get__ can come into play here.

[ilevkivskyi]

@orenbenkiki

Not sure what you mean

You can make MyAuto a descriptor (in stubs).

[orenbenkiki]

Ah, I see. This would require writing MyAuto[MyEnumClass] at the very least - and even then, I think mypy won't be smart enough to understand __get__ returns a MyEnumClass. Also, MyAuto is just one option - in other cases I initialize with an explicit value of some type (similar to explicitly initializing a standard enum value to an integer). In these cases, mypy thinks the property has the initializer's type instead of the enum's class type.

[elazarg]

I think it is a reasonable feature request to make instances of EnumMeta have enum-like types.

[emmatyping]

I feel that the underlying issue @orenbenkiki is hitting is #1020, since MyAuto returns a different type than it is. Although I'm not sure this is something that we should support.

Also, I'm not certain that we should support classes with EnumMeta as a metaclass, as it may not follow the assumptions we have made about Enum

[elazarg]

I wasn't talking about inheriting classes but classes that use it as metaclass; this is a declaration of intent just like inheriting Enum.

[emmatyping]

Ah, I mispoke, I mean classes that use it as a metaclass. Anyway, I don't see why we can't try that.

[elazarg]

I have implemented it - it's pretty easy. But it falsifies the assumption on typeshed that only Enum derivatives instantiate EnumClass, so typeshed should be fixed first.

[emmatyping]

@orenbenkiki I am going to close this, as the original ask is unlikely to happen, and EnumMeta is now checked properly. Feel free to leave a message if this is not satisfactory.

[elazarg]

@ethanhs What about #4319?

[emmatyping]

Oh, I thought that had been merged! My mistake.