[python] Message constructor argument enum typing #23670
Description
What version of protobuf and what language are you using?
Version: main
Language: Python
What did you do?
Steps to reproduce the behavior:
- Create a message with an enum
- Generate python and type stubs
- Create a message instance, and pass in an int as the enum value
What did you expect to see
No typing errors (works at runtime)
What did you see instead?
Message constructor is typed as a Union[<EnumType>, str] so I get an argument type error.
At runtime the message constructor accepts ints as enum arguments, but the typing does not reflect that.
test.proto:
syntax = "proto3";
package test;
enum TestEnum {
TEST_ENUM_UNSPECIFIED = 0;
TEST_ENUM_VALUE_1 = 1;
TEST_ENUM_VALUE_2 = 2;
}
message TestMessage {
string name = 1;
int32 id = 2;
TestEnum test_enum = 3;
}
test_pb2.pyi:
from google.protobuf.internal import enum_type_wrapper as _enum_type_wrapper
from google.protobuf import descriptor as _descriptor
from google.protobuf import message as _message
from typing import ClassVar as _ClassVar, Optional as _Optional, Union as _Union
DESCRIPTOR: _descriptor.FileDescriptor
class TestEnum(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
__slots__ = ()
TEST_ENUM_UNSPECIFIED: _ClassVar[TestEnum]
TEST_ENUM_VALUE_1: _ClassVar[TestEnum]
TEST_ENUM_VALUE_2: _ClassVar[TestEnum]
TEST_ENUM_UNSPECIFIED: TestEnum
TEST_ENUM_VALUE_1: TestEnum
TEST_ENUM_VALUE_2: TestEnum
class TestMessage(_message.Message):
__slots__ = ()
NAME_FIELD_NUMBER: _ClassVar[int]
ID_FIELD_NUMBER: _ClassVar[int]
TEST_ENUM_FIELD_NUMBER: _ClassVar[int]
name: str
id: int
test_enum: TestEnum
def __init__(self, name: _Optional[str] = ..., id: _Optional[int] = ..., test_enum: _Optional[_Union[TestEnum, str]] = ...) -> None: ...
Example code:
import test_pb2
test_pb2.TestMessage(test_enum=test_pb2.TestEnum.TEST_ENUM_VALUE_1)
test_pb2.TestMessage(test_enum="TEST_ENUM_VALUE_1")
test_pb2.TestMessage(test_enum=1)
The code runs fine, but mypy errors with: test.py:5: error: Argument "test_enum" to "TestMessage" has incompatible type "int"; expected "TestEnum | str | None" [arg-type]
I tested a change where I updated this line to be:
printer_->Print("_Union[$type_name$, str, int]", "type_name",
And now mypy passes