Generic versions of enum.Enum?

[bugreport1234]

(First of all, apologies if this is not the proper place for suggesting this, or if this has already been discussed, I couldn't find anything related to this so I opened an issue here.)

So I've recently come across a situation like this:
I've defined an enum

from enum import Enum

class MyEnum(Enum):
    a = 1
    b = 2
    c = 3

and need to use the value of an enum member later on:

some_func(MyEnum.a.value)

Analyzing this with mypy highlights the value of MyEnum.a.value as Any. That makes sense, as enum member values could be of any type.
However, in this case I know that all members of my enum should have int values, and no other type. In general, in most cases where I have used an enum the values have all been of a single type. I'd like to communicate this to the typing system somehow, but it doesn't seem possible with enum.Enum.

Changing the enum's type to enum.IntEnum lets mypy identify the value as int, however, using IntEnum is discouraged by the enum module documentation since it also makes enums comparable to other enums and to integers, which wouldn't actually be necessary for this use case. So as far as I understand using IntEnum wouldn't be ideal either.

Casting the .value to an int works, but it's more a workaround than a solution, and I hope this could somehow be done without casting.

To me, it seems this could be solved by introducing a generic Enum type in typing, similar to Sequence[T] and the other generics defined there.
With this, the code would look like

from typing import Enum

class MyEnum(Enum[int]):
    a = 1
    b = 2

some_func(MyEnum.a.value)

This could make the intent of the enum more clear, and would allow people to catch errors that would be introduced by defining an enum member with a different value type. It also would allow static checkers to infer the type of the enum member's value without having to use casts. Comparison operations would work the same way as for enum.Enum, and unlike enum.IntEnum.

If something like this could be considered for the typing module I'd be very grateful.

(Also, lastly, thanks for all the work on static typing in Python. It's helped me catch several bugs in my code so far, and I'm working on fully converting my project to make use of static typing, since it's been such a great help so far.)

[roganov]

+1 on this.

It's especially annoying that mypy requires to provide types of values explicitly in some cases, but later disregards this information and type of values becomes Any.
Example:

class Problematic(enum.Enum):

    a = frozenset()

mypy output:

 error: Need type annotation for 'a'

[JukkaL]

Changes to the typing module are best discussed at https://github.com/python/typing.

@roganov Can you file a separate issue about requiring a type annotation within the body of an enum?

[FuegoFro]

@JukkaL I'm a bit confused. It looks like this is that repository?

[ilevkivskyi]

Although this is potentially possible, I think this adds to much burden for little benefits. For most use cases, enums are just sets of unique constants. Also, type checkers could "remember" the type of .value, but again it seems to me the benefit is minimal as compared to the efforts required.

[TeamSpen210]

I don't think the generic class is needed. Enums are "frozen", you cannot add additional values after the class is constructed. They can't be subclassed either. MyPy could infer value as the union of all non-function class attributes. If a specific type annotation is desired, that could be applied to any member (likely the first).

[bugreport1234]

If type inference for enum attributes is possible, this could be an alternative solution to this problem. The main issue is that .value isn't assigned a type by the type checker, even though it could be. Type inference seems like a reasonable approach, and I understand that the generic class solution may be a bit too much compared to it. Type inference wouldn't require a change to existing code or the typing module, which would also make things easier.
While the benefit of type annotations for .value is fairly small, I believe there is enough that it's worth discussing. If enums could be completely typechecked, there'd be an additional reason to use enums instead of traditional "classes with constant attributes" enums.

[ilevkivskyi]

I think it is possible to implement this as a mypy plugin. But it is quite specialised, so it is unlikely that anyone from mypy core team will write the plugin. You of course can write the plugin yourself.

[kespindler]

A generic would be a cleaner solution, but this works:

from enum import Enum
from typing import TYPE_CHECKING


class Foo(Enum):
    value: int
    a = 1
    b = 2

if TYPE_CHECKING:
    reveal_type(Foo.a)
    reveal_type(Foo.a.value)

print(Foo.a)
print(Foo.a.value)
print(list(Foo))

[Photonios]

I sort of found a solution that works for me. My problem is that I want to define some utility methods on my enum base class. However, I wasn't able to type the return value of those methods properly.

First I tried this:

from enum import Enum
from typing import List, Tuple, TypeVar

T = TypeVar('T', bound=MyEnum)

class MyEnum(Enum):
    @classmethod  
    def choices(cls) -> List[Tuple[T, str]]:
          return [(item.value, item.name) for item in cls]

class Foo(MyEnum):
       A = 1
       B = 2

class Bar(MyEnum):
      A = "a"
      B = "b"

def do_stuff(foos: List[Tuple[Foo, str]]) -> None:
     print(foos)

do_stuff(Bar.choices()) # no error here!!

I tried making MyEnum generic, but MyPy throws an error that enums can't be generic. My solution was to make a mixin that is generic and overload the method in the mixin. This seems to work:

from enum import Enum
from typing import Any, List, Tuple, Generic, TypeVar

T = TypeVar('T')


class MyEnum(Enum):
    @classmethod
    def choices(cls) -> List[Tuple[Any, str]]:
          return [(item.value, item.name) for item in cls]


class MyEnumTypingMixin(Generic[T]):
    @classmethod
    def choices(cls) -> List[Tuple[T, str]]:
        ...

class Foo(MyEnumTypingMixin[Foo], MyEnum):
       A = 1
       B = 2

class Bar(MyEnumTypingMixin[Bar], MyEnum):
      A = "a"
      B = "b"

def do_stuff(foos: List[Tuple[Foo, str]]) -> None:
     print(foos)

do_stuff(Bar.choices()) # error: Argument 1 to "do_stuff" has incompatible type "List[Tuple[Bar, str]]"; expected "List[Tuple[Foo, str]]"

Admittedly this is a hack and I was surprised it works. Maybe it shouldn't? Someone with more experience in this might be able to tell better.

[sproshev]

@Photonios What if cls will be type hinted as Type[T]?

[ilevkivskyi]

@Photonios your use case doesn't require generic enums, as suggested above you can simply write:

from enum import Enum
from typing import Type, List, Tuple, TypeVar

T = TypeVar('T')

class MyEnum(Enum):
    @classmethod
    def choices(cls: Type[T]) -> List[Tuple[T, str]]:
          return [(item, item.name) for item in cls]  # error here is a mypy bug, ignore it

class Foo(MyEnum):
    A = 1
    B = 2

reveal_type(Foo.choices())  # Revealed type is 'builtins.list[Tuple[test.Foo*, builtins.str]]'

Also your original example type-checks because Any is not object.

Also it looks like there is a bug in your code (uncaught because of Any), IIUC you want (item, item.name), not (item.value, item.name). But if you actually want the latter, then yes, this requires generic enums, so that Foo will be Enum[int] and Bar will be Enum[str].

[pikeas]

Just bit by this. I was very surprised to find that .value() on an enum is typed as Any rather than resolved by the type system.

[srittau]

I believe that it would make sense to make Enum generic, at least in the stubs. But currently it would most likely disrupt too much. We would need something like #307 (defaults for generic arguments) to keep the disruption minimal.

[JelleZijlstra]

Ideally a type checker should be able to infer the generic type from the type of the values inside the enum. But that's not something we can do in typeshed.

[Ronnocerman]

The above solution to declare "value" as having "int" type did not work for me, as my type checker was noticing that "value" never got initialized.

What worked for me was declaring a "value" property and calling super.

from enum import Enum
from typing import Tuple

class MyEnum(Enum):
  my_val: Tuple[str, str] = ("foo", "bar")
  @property
  def value(self) -> Tuple[str, str]:
    return super().value

reveal_type(MyEnum.my_val.value)