diff --git a/third_party/2and3/google/protobuf/internal/enum_type_wrapper.pyi b/third_party/2and3/google/protobuf/internal/enum_type_wrapper.pyi index 174002d7a..1f26f3f8e 100644 --- a/third_party/2and3/google/protobuf/internal/enum_type_wrapper.pyi +++ b/third_party/2and3/google/protobuf/internal/enum_type_wrapper.pyi @@ -1,10 +1,19 @@ -from typing import Any, List, Tuple +from typing import Any, Generic, List, Tuple, TypeVar -class EnumTypeWrapper(object): - def __init__(self, enum_type: Any) -> None: ... - def Name(self, number: int) -> bytes: ... - def Value(self, name: bytes) -> int: ... - def keys(self) -> List[bytes]: ... - def values(self) -> List[int]: ... - @classmethod - def items(cls) -> List[Tuple[bytes, int]]: ... +from google.protobuf.descriptor import EnumDescriptor + +_V = TypeVar("_V", bound=int) + +# Expose a generic version so that those using mypy-protobuf +# can get autogenerated NewType wrapper around the int values +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]]: ... + +class EnumTypeWrapper(_EnumTypeWrapper[int]): ...