Use a single extension message for all the mypy_protobuf field extension types (#398)

This allows us to only expend a single extension number

Author: nipunn1313

CHANGELOG.md

@@ -1,5 +1,9 @@
 ## Upcoming
 
+- Prefer (mypy_protobuf.options).casttype to (mypy_protobuf.casttype)
+    - Allows us to use a single extension number
+    - Applies to casttype,keytype,valuetype
+    - Deprecate (but still support) the old-style extension
 - Support emitting module docstrings
 - Make generated code flake8 compatible
 - Convert extensions.proto to proto3

README.md

@@ -151,10 +151,10 @@ WORLD: MyEnum.ValueType  # 1
 M.proto
 ```proto
 message M {
-  uint32 user_id = 1 [(mypy_protobuf.casttype)="mymod.UserId"];
+  uint32 user_id = 1 [(mypy_protobuf.options).casttype="mymod.UserId"];
   map<uint32, string> email_by_uid = 2 [
-    (mypy_protobuf.keytype)="path/to/mymod.UserId",
-    (mypy_protobuf.valuetype)="path/to/mymod.Email"
+    (mypy_protobuf.options).keytype="path/to/mymod.UserId",
+    (mypy_protobuf.options).valuetype="path/to/mymod.Email"
   ];
 }
 ```

mypy_protobuf/extensions_pb2.py

@@ -6,15 +6,41 @@ import builtins
 import google.protobuf.descriptor
 import google.protobuf.descriptor_pb2
 import google.protobuf.internal.extension_dict
+import google.protobuf.message
+import typing_extensions
 
 DESCRIPTOR: google.protobuf.descriptor.FileDescriptor
 
+class FieldOptions(google.protobuf.message.Message):
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    CASTTYPE_FIELD_NUMBER: builtins.int
+    KEYTYPE_FIELD_NUMBER: builtins.int
+    VALUETYPE_FIELD_NUMBER: builtins.int
+    casttype: builtins.str
+    """Tells mypy-protobuf to use a specific newtype rather than the normal type for this field."""
+    keytype: builtins.str
+    """Tells mypy-protobuf to use a specific type for keys; only makes sense on map fields"""
+    valuetype: builtins.str
+    """Tells mypy-protobuf to use a specific type for values; only makes sense on map fields"""
+    def __init__(
+        self,
+        *,
+        casttype: builtins.str = ...,
+        keytype: builtins.str = ...,
+        valuetype: builtins.str = ...,
+    ) -> None: ...
+    def ClearField(self, field_name: typing_extensions.Literal["casttype", b"casttype", "keytype", b"keytype", "valuetype", b"valuetype"]) -> None: ...
+
+global___FieldOptions = FieldOptions
+
+OPTIONS_FIELD_NUMBER: builtins.int
 CASTTYPE_FIELD_NUMBER: builtins.int
 KEYTYPE_FIELD_NUMBER: builtins.int
 VALUETYPE_FIELD_NUMBER: builtins.int
+options: google.protobuf.internal.extension_dict._ExtensionFieldDescriptor[google.protobuf.descriptor_pb2.FieldOptions, global___FieldOptions]
+"""Custom field options from mypy-protobuf"""
 casttype: google.protobuf.internal.extension_dict._ExtensionFieldDescriptor[google.protobuf.descriptor_pb2.FieldOptions, builtins.str]
-"""Tells mypy to use a specific newtype rather than the normal type for this field."""
+"""Legacy fields. Prefer to use ones within `options` instead."""
 keytype: google.protobuf.internal.extension_dict._ExtensionFieldDescriptor[google.protobuf.descriptor_pb2.FieldOptions, builtins.str]
-"""Tells mypy to use a specific type for keys; only makes sense on map fields"""
 valuetype: google.protobuf.internal.extension_dict._ExtensionFieldDescriptor[google.protobuf.descriptor_pb2.FieldOptions, builtins.str]
-"""Tells mypy to use a specific type for values; only makes sense on map fields"""

mypy_protobuf/extensions_pb2.pyi

@@ -634,10 +634,18 @@ def _map_key_value_types(
         key_field: d.FieldDescriptorProto,
         value_field: d.FieldDescriptorProto,
     ) -> Tuple[str, str]:
-        key_casttype = map_field.options.Extensions[extensions_pb2.keytype]
+        oldstyle_keytype = map_field.options.Extensions[extensions_pb2.keytype]
+        if oldstyle_keytype:
+            print(f"Warning: Map Field {map_field.name}: (mypy_protobuf.keytype) is deprecated. Prefer (mypy_protobuf.options).keytype", file=sys.stderr)
+        key_casttype = map_field.options.Extensions[extensions_pb2.options].keytype or oldstyle_keytype
         ktype = self._import_casttype(key_casttype) if key_casttype else self.python_type(key_field)
-        value_casttype = map_field.options.Extensions[extensions_pb2.valuetype]
+
+        oldstyle_valuetype = map_field.options.Extensions[extensions_pb2.valuetype]
+        if oldstyle_valuetype:
+            print(f"Warning: Map Field {map_field.name}: (mypy_protobuf.valuetype) is deprecated. Prefer (mypy_protobuf.options).valuetype", file=sys.stderr)
+        value_casttype = map_field.options.Extensions[extensions_pb2.options].valuetype or map_field.options.Extensions[extensions_pb2.valuetype]
         vtype = self._import_casttype(value_casttype) if value_casttype else self.python_type(value_field)
+
         return ktype, vtype
 
     def _callable_type(self, method: d.MethodDescriptorProto) -> str:
@@ -758,7 +766,10 @@ def python_type(self, field: d.FieldDescriptorProto, generic_container: bool = F
           - Mapping[k, v] rather than MessageMap[k, v]
           Can be useful for input types (eg constructor)
         """
-        casttype = field.options.Extensions[extensions_pb2.casttype]
+        oldstyle_casttype = field.options.Extensions[extensions_pb2.casttype]
+        if oldstyle_casttype:
+            print(f"Warning: Field {field.name}: (mypy_protobuf.casttype) is deprecated. Prefer (mypy_protobuf.options).casttype", file=sys.stderr)
+        casttype = field.options.Extensions[extensions_pb2.options].casttype or oldstyle_casttype
         if casttype:
             return self._import_casttype(casttype)
 

mypy_protobuf/main.py

@@ -4,12 +4,22 @@ package mypy_protobuf;
 
 import "google/protobuf/descriptor.proto";
 
+message FieldOptions {
+    // Tells mypy-protobuf to use a specific newtype rather than the normal type for this field.
+    string casttype = 1;
+    // Tells mypy-protobuf to use a specific type for keys; only makes sense on map fields
+    string keytype = 2;
+    // Tells mypy-protobuf to use a specific type for values; only makes sense on map fields
+    string valuetype = 3;
+
+}
+
 extend google.protobuf.FieldOptions {
-    // Tells mypy to use a specific newtype rather than the normal type for this field.
-    string casttype = 60000;
+    // Custom field options from mypy-protobuf
+    FieldOptions options = 60004;
 
-    // Tells mypy to use a specific type for keys; only makes sense on map fields
-    string keytype = 60002;
-    // Tells mypy to use a specific type for values; only makes sense on map fields
-    string valuetype = 60003;
+    // Legacy fields. Prefer to use ones within `options` instead.
+    string casttype = 60000 [deprecated = true];
+    string keytype = 60002 [deprecated = true];
+    string valuetype = 60003 [deprecated = true];
 }

proto/mypy_protobuf/extensions.proto

@@ -8,6 +8,8 @@ syntax = "proto3";
 /* package test3 */
 package test3;
 
+import "mypy_protobuf/extensions.proto";
+
 enum OuterEnum {
   UNKNOWN = 0;
   FOO3 = 1;
@@ -54,4 +56,11 @@ message SimpleProto3 {
   map<int32, OuterMessage3> map_message = 17;
 
   optional string an_optional_string = 18;
+
+  uint32 user_id = 19 [(mypy_protobuf.options).casttype="test/test_generated_mypy.UserId"];
+  string email = 20 [(mypy_protobuf.options).casttype="test/test_generated_mypy.Email"];
+  map<uint32, string> email_by_uid = 21 [
+    (mypy_protobuf.options).keytype="test/test_generated_mypy.UserId",
+    (mypy_protobuf.options).valuetype="test/test_generated_mypy.Email"
+  ];
 }

proto/testproto/test3.proto

@@ -170,7 +170,7 @@ for PY_VER in $PY_VER_UNIT_TESTS; do
             echo -e "${RED}test_negative/output.expected.$PY_VER_MYPY_TARGET didnt match. Copying over for you. Now rerun${NC}"
 
             # Copy over all the mypy results for the developer.
-            call_mypy 3.8.11 "${NEGATIVE_FILES[@]}"
+            call_mypy $PY_VER "${NEGATIVE_FILES[@]}"
             cp "$MYPY_OUTPUT/mypy_output" test_negative/output.expected.3.8
             cp "$MYPY_OUTPUT/mypy_output.omit_linenos" test_negative/output.expected.3.8.omit_linenos
             exit 1

run_test.sh

@@ -6,15 +6,41 @@ import builtins
 import google.protobuf.descriptor
 import google.protobuf.descriptor_pb2
 import google.protobuf.internal.extension_dict
+import google.protobuf.message
+import typing_extensions
 
 DESCRIPTOR: google.protobuf.descriptor.FileDescriptor
 
+class FieldOptions(google.protobuf.message.Message):
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    CASTTYPE_FIELD_NUMBER: builtins.int
+    KEYTYPE_FIELD_NUMBER: builtins.int
+    VALUETYPE_FIELD_NUMBER: builtins.int
+    casttype: builtins.str
+    """Tells mypy-protobuf to use a specific newtype rather than the normal type for this field."""
+    keytype: builtins.str
+    """Tells mypy-protobuf to use a specific type for keys; only makes sense on map fields"""
+    valuetype: builtins.str
+    """Tells mypy-protobuf to use a specific type for values; only makes sense on map fields"""
+    def __init__(
+        self,
+        *,
+        casttype: builtins.str = ...,
+        keytype: builtins.str = ...,
+        valuetype: builtins.str = ...,
+    ) -> None: ...
+    def ClearField(self, field_name: typing_extensions.Literal["casttype", b"casttype", "keytype", b"keytype", "valuetype", b"valuetype"]) -> None: ...
+
+global___FieldOptions = FieldOptions
+
+OPTIONS_FIELD_NUMBER: builtins.int
 CASTTYPE_FIELD_NUMBER: builtins.int
 KEYTYPE_FIELD_NUMBER: builtins.int
 VALUETYPE_FIELD_NUMBER: builtins.int
+options: google.protobuf.internal.extension_dict._ExtensionFieldDescriptor[google.protobuf.descriptor_pb2.FieldOptions, global___FieldOptions]
+"""Custom field options from mypy-protobuf"""
 casttype: google.protobuf.internal.extension_dict._ExtensionFieldDescriptor[google.protobuf.descriptor_pb2.FieldOptions, builtins.str]
-"""Tells mypy to use a specific newtype rather than the normal type for this field."""
+"""Legacy fields. Prefer to use ones within `options` instead."""
 keytype: google.protobuf.internal.extension_dict._ExtensionFieldDescriptor[google.protobuf.descriptor_pb2.FieldOptions, builtins.str]
-"""Tells mypy to use a specific type for keys; only makes sense on map fields"""
 valuetype: google.protobuf.internal.extension_dict._ExtensionFieldDescriptor[google.protobuf.descriptor_pb2.FieldOptions, builtins.str]
-"""Tells mypy to use a specific type for values; only makes sense on map fields"""

test/generated/mypy_protobuf/extensions_pb2.pyi

@@ -8,6 +8,7 @@ import google.protobuf.descriptor
 import google.protobuf.internal.containers
 import google.protobuf.internal.enum_type_wrapper
 import google.protobuf.message
+import test.test_generated_mypy
 import typing
 import typing_extensions
 
@@ -92,6 +93,21 @@ class SimpleProto3(google.protobuf.message.Message):
         def HasField(self, field_name: typing_extensions.Literal["value", b"value"]) -> builtins.bool: ...
         def ClearField(self, field_name: typing_extensions.Literal["key", b"key", "value", b"value"]) -> None: ...
 
+    class EmailByUidEntry(google.protobuf.message.Message):
+        DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+        KEY_FIELD_NUMBER: builtins.int
+        VALUE_FIELD_NUMBER: builtins.int
+        key: builtins.int
+        value: builtins.str
+        def __init__(
+            self,
+            *,
+            key: builtins.int = ...,
+            value: builtins.str = ...,
+        ) -> None: ...
+        def ClearField(self, field_name: typing_extensions.Literal["key", b"key", "value", b"value"]) -> None: ...
+
     A_STRING_FIELD_NUMBER: builtins.int
     A_REPEATED_STRING_FIELD_NUMBER: builtins.int
     A_OUTER_ENUM_FIELD_NUMBER: builtins.int
@@ -110,6 +126,9 @@ class SimpleProto3(google.protobuf.message.Message):
     MAP_SCALAR_FIELD_NUMBER: builtins.int
     MAP_MESSAGE_FIELD_NUMBER: builtins.int
     AN_OPTIONAL_STRING_FIELD_NUMBER: builtins.int
+    USER_ID_FIELD_NUMBER: builtins.int
+    EMAIL_FIELD_NUMBER: builtins.int
+    EMAIL_BY_UID_FIELD_NUMBER: builtins.int
     a_string: builtins.str
     @property
     def a_repeated_string(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]: ...
@@ -137,6 +156,10 @@ class SimpleProto3(google.protobuf.message.Message):
     @property
     def map_message(self) -> google.protobuf.internal.containers.MessageMap[builtins.int, global___OuterMessage3]: ...
     an_optional_string: builtins.str
+    user_id: test.test_generated_mypy.UserId
+    email: test.test_generated_mypy.Email
+    @property
+    def email_by_uid(self) -> google.protobuf.internal.containers.ScalarMap[test.test_generated_mypy.UserId, test.test_generated_mypy.Email]: ...
     def __init__(
         self,
         *,
@@ -158,9 +181,12 @@ class SimpleProto3(google.protobuf.message.Message):
         map_scalar: collections.abc.Mapping[builtins.int, builtins.str] | None = ...,
         map_message: collections.abc.Mapping[builtins.int, global___OuterMessage3] | None = ...,
         an_optional_string: builtins.str | None = ...,
+        user_id: test.test_generated_mypy.UserId = ...,
+        email: test.test_generated_mypy.Email = ...,
+        email_by_uid: collections.abc.Mapping[test.test_generated_mypy.UserId, test.test_generated_mypy.Email] | None = ...,
     ) -> None: ...
     def HasField(self, field_name: typing_extensions.Literal["OuterMessage3", b"OuterMessage3", "_an_optional_string", b"_an_optional_string", "a_oneof", b"a_oneof", "a_oneof_1", b"a_oneof_1", "a_oneof_2", b"a_oneof_2", "an_optional_string", b"an_optional_string", "b_oneof", b"b_oneof", "b_oneof_1", b"b_oneof_1", "b_oneof_2", b"b_oneof_2", "bool", b"bool", "inner_enum_in_oneof", b"inner_enum_in_oneof", "outer_enum_in_oneof", b"outer_enum_in_oneof", "outer_message", b"outer_message", "outer_message_in_oneof", b"outer_message_in_oneof"]) -> builtins.bool: ...
-    def ClearField(self, field_name: typing_extensions.Literal["OuterEnum", b"OuterEnum", "OuterMessage3", b"OuterMessage3", "_an_optional_string", b"_an_optional_string", "a_oneof", b"a_oneof", "a_oneof_1", b"a_oneof_1", "a_oneof_2", b"a_oneof_2", "a_outer_enum", b"a_outer_enum", "a_repeated_string", b"a_repeated_string", "a_string", b"a_string", "an_optional_string", b"an_optional_string", "b_oneof", b"b_oneof", "b_oneof_1", b"b_oneof_1", "b_oneof_2", b"b_oneof_2", "bool", b"bool", "inner_enum", b"inner_enum", "inner_enum_in_oneof", b"inner_enum_in_oneof", "map_message", b"map_message", "map_scalar", b"map_scalar", "outer_enum_in_oneof", b"outer_enum_in_oneof", "outer_message", b"outer_message", "outer_message_in_oneof", b"outer_message_in_oneof"]) -> None: ...
+    def ClearField(self, field_name: typing_extensions.Literal["OuterEnum", b"OuterEnum", "OuterMessage3", b"OuterMessage3", "_an_optional_string", b"_an_optional_string", "a_oneof", b"a_oneof", "a_oneof_1", b"a_oneof_1", "a_oneof_2", b"a_oneof_2", "a_outer_enum", b"a_outer_enum", "a_repeated_string", b"a_repeated_string", "a_string", b"a_string", "an_optional_string", b"an_optional_string", "b_oneof", b"b_oneof", "b_oneof_1", b"b_oneof_1", "b_oneof_2", b"b_oneof_2", "bool", b"bool", "email", b"email", "email_by_uid", b"email_by_uid", "inner_enum", b"inner_enum", "inner_enum_in_oneof", b"inner_enum_in_oneof", "map_message", b"map_message", "map_scalar", b"map_scalar", "outer_enum_in_oneof", b"outer_enum_in_oneof", "outer_message", b"outer_message", "outer_message_in_oneof", b"outer_message_in_oneof", "user_id", b"user_id"]) -> None: ...
     @typing.overload
     def WhichOneof(self, oneof_group: typing_extensions.Literal["_an_optional_string", b"_an_optional_string"]) -> typing_extensions.Literal["an_optional_string"] | None: ...
     @typing.overload

test/generated/testproto/test3_pb2.pyi

@@ -16,6 +16,7 @@
 import testproto.test_pb2 as test_pb2
 from google.protobuf.descriptor import FieldDescriptor
 from google.protobuf.internal import api_implementation
+from google.protobuf.internal.containers import ScalarMap
 from google.protobuf.message import Message
 from testproto.Capitalized.Capitalized_pb2 import Upper, lower, lower2
 from testproto.reexport_pb2 import FOO3 as ReexportedFOO3
@@ -41,6 +42,7 @@
     Simple1,
     Simple2,
 )
+from typing_extensions import assert_type
 
 # C++ python API implementation has some semantic differences from pure python
 # We mainly focus on testing the C++ impl - but just have some flags here
@@ -474,12 +476,27 @@ def test_mapping_type() -> None:
 
 def test_casttype() -> None:
     s = Simple1()
+    assert_type(s.user_id, UserId)
+    assert_type(s.email, Email)
+    assert_type(s.email_by_uid, ScalarMap[UserId, Email])
+
     s.user_id = UserId(33)
     assert s.user_id == 33
     s.email = Email("[email protected]")
     assert s.email == "[email protected]"
     s.email_by_uid[UserId(33)] = Email("[email protected]")
 
+    s2 = SimpleProto3()
+    assert_type(s2.user_id, UserId)
+    assert_type(s2.email, Email)
+    assert_type(s2.email_by_uid, ScalarMap[UserId, Email])
+
+    s2.user_id = UserId(33)
+    assert s2.user_id == 33
+    s2.email = Email("[email protected]")
+    assert s2.email == "[email protected]"
+    s2.email_by_uid[UserId(33)] = Email("[email protected]")
+
 
 def test_reserved_keywords() -> None:
     with pytest.raises(AttributeError, match="module.*has no attribute 'asdf'"):