[python] Message constructor argument enum typing

What version of protobuf and what language are you using? Version: main Language: Python

What did you do? Steps to reproduce the behavior:

1. Create a message with an enum
2. Generate python and type stubs
3. 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