Support repeated extensions

nipunn1313 · Jul 16, 2021 · f51572b · f51572b
1 parent 4f20127
commit f51572b
Show file tree

Hide file tree

Showing 8 changed files with 63 additions and 17 deletions.
diff --git a/proto/testproto/test3.proto b/proto/testproto/test3.proto
@@ -9,7 +9,7 @@ enum OuterEnum {
 }
 
 message OuterMessage3 {
-  bool a_bool = 1;
+  string a_string = 1;
 }
 
 message SimpleProto3 {

diff --git a/proto/testproto/test_extensions3.proto b/proto/testproto/test_extensions3.proto
@@ -1,17 +1,34 @@
 syntax = "proto3";
 
 import "google/protobuf/descriptor.proto";
+import "testproto/test3.proto";
 
-package test;
+package test3;
 
 extend google.protobuf.FieldOptions {
   string test_field_extension = 50000;
 }
 
 extend google.protobuf.MessageOptions {
-  string test_message_option = 51234;
+  string scalar_option = 51234;
+  repeated string repeated_scalar_option = 51235;
+  OuterEnum enum_option = 51236;
+  repeated OuterEnum repeated_enum_option = 51237;
+  OuterMessage3 msg_option = 51238;
+  repeated OuterMessage3 repeated_msg_option = 51239;
 }
 
 message MessageOptionsTestMsg {
-    option (test_message_option) = "Hello world!";
+    option (scalar_option) = "Hello world!";
+    option (repeated_scalar_option) = "A";
+    option (repeated_scalar_option) = "B";
+    option (repeated_scalar_option) = "C";
+
+    option (enum_option) = FOO3;
+    option (repeated_enum_option) = FOO3;
+    option (repeated_enum_option) = BAR3;
+
+    option (msg_option).a_string = "Hello OuterMessage3";
+    option (repeated_msg_option) = {a_string: "Hello OuterMessage3 A"};
+    option (repeated_msg_option) = {a_string: "Hello OuterMessage3 B"};
 }
diff --git a/test/generated/testproto/test3_pb2.pyi.expected b/test/generated/testproto/test3_pb2.pyi.expected
@@ -29,14 +29,14 @@ class _OuterEnum(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[Out
 
 class OuterMessage3(google.protobuf.message.Message):
     DESCRIPTOR: google.protobuf.descriptor.Descriptor = ...
-    A_BOOL_FIELD_NUMBER: builtins.int
-    a_bool: builtins.bool = ...
+    A_STRING_FIELD_NUMBER: builtins.int
+    a_string: typing.Text = ...
 
     def __init__(self,
         *,
-        a_bool : builtins.bool = ...,
+        a_string : typing.Text = ...,
         ) -> None: ...
-    def ClearField(self, field_name: typing_extensions.Literal[u"a_bool",b"a_bool"]) -> None: ...
+    def ClearField(self, field_name: typing_extensions.Literal[u"a_string",b"a_string"]) -> None: ...
 global___OuterMessage3 = OuterMessage3
 
 class SimpleProto3(google.protobuf.message.Message):

diff --git a/test/generated/testproto/test_extensions3_pb2.pyi.expected b/test/generated/testproto/test_extensions3_pb2.pyi.expected
@@ -4,8 +4,10 @@ isort:skip_file
 """
 import google.protobuf.descriptor
 import google.protobuf.descriptor_pb2
+import google.protobuf.internal.containers
 import google.protobuf.internal.extension_dict
 import google.protobuf.message
+import testproto.test3_pb2
 import typing
 
 DESCRIPTOR: google.protobuf.descriptor.FileDescriptor = ...
@@ -19,4 +21,14 @@ global___MessageOptionsTestMsg = MessageOptionsTestMsg
 
 test_field_extension: google.protobuf.internal.extension_dict._ExtensionFieldDescriptor[google.protobuf.descriptor_pb2.FieldOptions, typing.Text] = ...
 
-test_message_option: google.protobuf.internal.extension_dict._ExtensionFieldDescriptor[google.protobuf.descriptor_pb2.MessageOptions, typing.Text] = ...
+scalar_option: google.protobuf.internal.extension_dict._ExtensionFieldDescriptor[google.protobuf.descriptor_pb2.MessageOptions, typing.Text] = ...
+
+repeated_scalar_option: google.protobuf.internal.extension_dict._ExtensionFieldDescriptor[google.protobuf.descriptor_pb2.MessageOptions, google.protobuf.internal.containers.RepeatedScalarFieldContainer[typing.Text]] = ...
+
+enum_option: google.protobuf.internal.extension_dict._ExtensionFieldDescriptor[google.protobuf.descriptor_pb2.MessageOptions, testproto.test3_pb2.OuterEnum.V] = ...
+
+repeated_enum_option: google.protobuf.internal.extension_dict._ExtensionFieldDescriptor[google.protobuf.descriptor_pb2.MessageOptions, google.protobuf.internal.containers.RepeatedScalarFieldContainer[testproto.test3_pb2.OuterEnum.V]] = ...
+
+msg_option: google.protobuf.internal.extension_dict._ExtensionFieldDescriptor[google.protobuf.descriptor_pb2.MessageOptions, testproto.test3_pb2.OuterMessage3] = ...
+
+repeated_msg_option: google.protobuf.internal.extension_dict._ExtensionFieldDescriptor[google.protobuf.descriptor_pb2.MessageOptions, google.protobuf.internal.containers.RepeatedCompositeFieldContainer[testproto.test3_pb2.OuterMessage3]] = ...
diff --git a/test/test_generated_mypy.py b/test/test_generated_mypy.py
@@ -38,11 +38,20 @@
     Simple2,
 )
 from testproto.test3_pb2 import (
+    BAR3,
     FOO3,
     OuterMessage3,
     SimpleProto3,
 )
-from testproto.test_extensions3_pb2 import MessageOptionsTestMsg, test_message_option
+from testproto.test_extensions3_pb2 import (
+    MessageOptionsTestMsg,
+    scalar_option,
+    repeated_scalar_option,
+    enum_option,
+    repeated_enum_option,
+    msg_option,
+    repeated_msg_option,
+)
 from testproto.Capitalized.Capitalized_pb2 import lower, lower2, Upper
 
 from typing import (
@@ -451,9 +460,17 @@ def test_extensions_proto2():
 def test_extensions_proto3():
     # type: () -> None
     assert (
-        MessageOptionsTestMsg.DESCRIPTOR.GetOptions().Extensions[test_message_option]
+        MessageOptionsTestMsg.DESCRIPTOR.GetOptions().Extensions[scalar_option]
         == "Hello world!"
     )
+    assert MessageOptionsTestMsg.DESCRIPTOR.GetOptions().Extensions[repeated_scalar_option] == ["A", "B", "C"]
+    assert MessageOptionsTestMsg.DESCRIPTOR.GetOptions().Extensions[enum_option] == FOO3
+    assert MessageOptionsTestMsg.DESCRIPTOR.GetOptions().Extensions[repeated_enum_option] == [FOO3, BAR3]
+    assert MessageOptionsTestMsg.DESCRIPTOR.GetOptions().Extensions[msg_option] == OuterMessage3(a_string="Hello OuterMessage3")
+    assert list(MessageOptionsTestMsg.DESCRIPTOR.GetOptions().Extensions[repeated_msg_option]) == [
+        OuterMessage3(a_string="Hello OuterMessage3 A"),
+        OuterMessage3(a_string="Hello OuterMessage3 B"),
+    ]
 
 
 def test_constructor_proto2():
@@ -499,15 +516,15 @@ def test_mapping_type():
     s.map_scalar[5] = "abcd"
     assert s.map_scalar[5] == "abcd"
 
-    s.map_message[5].a_bool = True
-    assert s.map_message[5] == OuterMessage3(a_bool=True)
+    s.map_message[5].a_string = "hi"
+    assert s.map_message[5] == OuterMessage3(a_string="hi")
 
     assert s.map_message.get_or_create(6) == OuterMessage3()
     assert s.map_message[6] == OuterMessage3()
     assert s.map_message.get_or_create(6) == OuterMessage3()
 
     s2 = SimpleProto3(
-        map_scalar={5: "abcd"}, map_message={5: OuterMessage3(a_bool=True)}
+        map_scalar={5: "abcd"}, map_message={5: OuterMessage3(a_string="hi")}
     )
 
 

diff --git a/test_negative/negative.py b/test_negative/negative.py
@@ -163,7 +163,7 @@
 
 # In proto2 - you can pass in None for primitive, but not in proto3
 Simple2(a_string=None)
-OuterMessage3(a_bool=None)  # E:2.7 E:3.8
+OuterMessage3(a_string=None)  # E:2.7 E:3.8
 
 # Repeated scalar fields are not assignable only extendable
 s9 = Simple1()

diff --git a/test_negative/output.expected.2.7 b/test_negative/output.expected.2.7
@@ -67,7 +67,7 @@ test_negative/negative.py:159: error: Argument "user_id" to "Simple1" has incomp
 test_negative/negative.py:159: error: Argument "email" to "Simple1" has incompatible type "str"; expected "Optional[Email]"
 test_negative/negative.py:159: error: Dict entry 0 has incompatible type "int": "str"; expected "UserId": "Email"
 test_negative/negative.py:162: error: Module "testproto.reexport_pb2" has no attribute "Inner"
-test_negative/negative.py:166: error: Argument "a_bool" to "OuterMessage3" has incompatible type "None"; expected "bool"
+test_negative/negative.py:166: error: Argument "a_string" to "OuterMessage3" has incompatible type "None"; expected "unicode"
 test_negative/negative.py:171: error: Property "a_repeated_string" defined in "Simple1" is read-only
 test_negative/negative.py:172: error: Property "rep_inner_enum" defined in "Simple1" is read-only
 Found 62 errors in 1 file (checked 17 source files)
diff --git a/test_negative/output.expected.3.8 b/test_negative/output.expected.3.8
@@ -89,7 +89,7 @@ test_negative/negative.py:159: error: Argument "user_id" to "Simple1" has incomp
 test_negative/negative.py:159: error: Argument "email" to "Simple1" has incompatible type "str"; expected "Optional[Email]"
 test_negative/negative.py:159: error: Dict entry 0 has incompatible type "int": "str"; expected "UserId": "Email"
 test_negative/negative.py:162: error: Module "testproto.reexport_pb2" has no attribute "Inner"
-test_negative/negative.py:166: error: Argument "a_bool" to "OuterMessage3" has incompatible type "None"; expected "bool"
+test_negative/negative.py:166: error: Argument "a_string" to "OuterMessage3" has incompatible type "None"; expected "str"
 test_negative/negative.py:171: error: Property "a_repeated_string" defined in "Simple1" is read-only
 test_negative/negative.py:172: error: Property "rep_inner_enum" defined in "Simple1" is read-only
 Found 78 errors in 2 files (checked 32 source files)