Better mypy typing for enum values in python

[nipunn1313]

Hi - I'm the maintainer of https://github.com/dropbox/mypy-protobuf - a protoc plugin for autogenerating mypy stubs compatible with the built in python generator.

What language does this apply to?
Python

Describe the problem you are trying to solve.
Generated enum values are pure ints (per the spec). We prefer type them with a NewType wrapper type around int - specific to the enum. See the typeshed stub
https://github.com/python/typeshed/blob/master/third_party/2and3/google/protobuf/internal/enum_type_wrapper.pyi#L9

class _EnumTypeWrapper(Generic[_V]):
    DESCRIPTOR: EnumDescriptor
    def __init__(self, enum_type: EnumDescriptor) -> None: ...
    def Name(self, number: _V) -> str: ...
    def Value(self, name: str) -> _V: ...
    def keys(self) -> List[str]: ...
    def values(self) -> List[_V]: ...
    def items(self) -> List[Tuple[str, _V]]: ...

I've been cooking up a strategy inside mypy-protobuf where a proto

enum MyEnum {
   FOO = 0;
   BAR = 1;
}

Autogenerates pyi stubs

class _MyEnum(google.protobuf.internal.enum_type_wrapper.EnumTypeWrapper[OuterEnum.V], builtins.type):
    DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor = ...
    FOO = OuterEnum.V(0)
    BAR = OuterEnum.V(1)

class MyEnum(metaclass=_MyEnum):
    V = typing.NewType('V', builtins.int)

This allows callsites to type their code as

x: MyEnum.V = MyEnum.Value("FOO")

Unfortunately - because .V does not exist at runtime (only in the pyi stubs) - this does not work.
Fortunately, you can work around this with:

x: "MyEnum.V" = MyEnum.Value("FOO")

Describe the solution you'd like
I'd like to propose adding V = int runtime type alias to the EnumTypeWrapper (in enum_type_wrapper.py) - to avoid the need for this workaround.

class EnumTypeWrapper(object):
  """A utility for finding the names of enum values."""
  # This is a type alias, which mypy typing stubs can type as
  # a genericized parameter constrained to an int, allowing subclasses
  # to be typed with more constraint
  # Eg.
  # def Name(self, number: MyGeneratedEnum.V) -> str
  V = int
  ... [rest of impl omitted]

This would mean that enum variants named "V" would need to jump through a hoop using getattr to be accessed, but this is already the case with "Name", "Value", "keys", etc - so it seems feasible as an extension. I'm open to other naming ideas for this.

Could do something like this for more clarity - though it does add a dependency on typing module for python <= 3.5

from typing import NewType
V = NewType('V', int)

I would also be open to discussing other strategies we could employ across mypy-protobuf, protoc, and typeshed in order to get better mypy typing for enum values (better than just typing as an int) - other than this V inner type alias.

Describe alternatives you've considered
I've also tried mypy typing V at a sibling scope to MyEnum (eg MyEnumValue = NewType('MyEnumValue', int), but this is no improvement - as you still need to play tricks to import it from real code.

Thanks friends! Happy to open discussion.

[nipunn1313]

Another idea we've tried - would be to make MyEnum an instance of int. Then we would be typing the enumtypewrapper instance as MyEnum as well as the int values as MyEnum.

This seems a bit of a departure from reality of the generated code - as those are actually two separate types (instance of EnumTypeWrapper vs instance of int), but it can be made workable. The typing would allow folks to accidentally try and call enumtypewrapper methods on the ints themselves, but it would cut down on this "extra" type.

I'd softly lean toward the V solution as it makes clear in the typing that the EnumTypeWrapper and the int values are distinct objects of distinct types.

[nipunn1313]

Hi! Any updated thoughts on #8182. I'm planning on going forward with the V solution in mypy-protobuf - and it would be ergonomic for calling code to have #8182 to write code like this!

def f(x: MyEnum.V):
    print(x)

f(MyEnum.Value("FOO"))

rather than like this

def f(x: 'MyEnum.V'):
    print(x)

f(MyEnum.Value("FOO"))

[nipunn1313]

I realize - that with python 3.10 (or w/ from __future__ import annotations) - python will no longer try to execute the type annotations, meaning that this sort of change may no longer be necessary.

[nipunn1313]

Actually realized - even with from __future__ import annotations - it is still helpful to be able to inline code that casts to the type wrapper. Eg

x = MyEnum.V(5)

[nipunn1313]

Hi @anandolee - wanted to check in on this task and #8182

[nipunn1313]

Hey @anandolee and/or @acozzette - wanted to check in here to see if there are next steps I can help with. Happy to put in a bit of work here (either technically / documentation / derisking) if that would be helpful - just point me to what would be useful here.

If there's some high level concern, lmk as that might be helpful - either for coming up with a different strategy, or to communicate to folks being bothered by the status quo.

[jaywhy13]

@nipunn1313, I stumbled across this issue trying to understand why MyEnum.FOO didn't typecheck. I'm reviewing your suggestion and trying to understand how you settled on using V? Was it not possible to define a static class that had the enum attributes?

[nipunn1313]

Hi @jaywhy13.
protobuf generates an EnumTypeWrapper (static class) with helper methods, and a raw int for the value type.

Rather than typing as a raw int, mypy-protobuf wanted stronger typing - specific to the enum itself - a pyi file like this:

class MyEnum(EnumTypeWrapper):
  V = NewType('V', int)

  FOO: int

I suppose you're wondering why we didn't generate typing stubs more like

class MyEnum(EnumTypeWrapper):
   class V(IntEnum):
       FOO: int

This could work as well, but still ends up requiring a V. May end up being somewhat misleading as the actual generated protobuf code uses an int and not an IntEnum. NewType around int is truer to the generated protobuf code.

In either case, the base code in protobuf itself would still need something like V = int

In selecting V - chose it to be short for Value - but also easy to type (just one letter). Most places in the code would show something like MyEnum.V and be clear.

[nipunn1313]

For clarity - I worked on mypy-protobuf and not on protobuf itself, so I don't currently have the flexibility to make changes to both. I'm proposing a minimal change to protobuf to the core protobuf team in #8182 - so that mypy-protobuf can work better.