Fix module level extensions (#245)

Update types-protobuf pin from typeshed

Depends on python/typeshed#5774
Fixes #244

CHANGELOG.md

@@ -1,5 +1,7 @@
 ## Upcoming
 
+- Fix #244 - support extensions defined at module scope with proper types.
+
 ## 2.6
 
 - Bump protoc support to 3.17.3

README.md

@@ -16,7 +16,7 @@ See [Changelog](CHANGELOG.md) for recent changes.
 ## Requirements to run mypy on stubs generated by mypy-protobuf
 [mypy >= v0.910](https://pypi.org/project/mypy)
 [python-protobuf >= 3.17.3](https://pypi.org/project/protobuf/) - matching protoc release
-[types-protobuf==0.1.14](https://pypi.org/project/types-protobuf/) - for stubs from the google.protobuf library
+[types-protobuf==3.17.1](https://pypi.org/project/types-protobuf/) - for stubs from the google.protobuf library
 
 ### To run mypy on code generated with grpc plugin - you'll additionally need
 [grpcio>=1.38.1](https://pypi.org/project/grpcio/)

mypy_protobuf/extensions_pb2.pyi

@@ -3,11 +3,14 @@
 isort:skip_file
 """
 import google.protobuf.descriptor
+import google.protobuf.descriptor_pb2
+import google.protobuf.internal.extension_dict
+import typing
 
 DESCRIPTOR: google.protobuf.descriptor.FileDescriptor = ...
 
-casttype: google.protobuf.descriptor.FieldDescriptor = ...
+casttype: google.protobuf.internal.extension_dict._ExtensionFieldDescriptor[google.protobuf.descriptor_pb2.FieldOptions, typing.Text] = ...
 
-keytype: google.protobuf.descriptor.FieldDescriptor = ...
+keytype: google.protobuf.internal.extension_dict._ExtensionFieldDescriptor[google.protobuf.descriptor_pb2.FieldOptions, typing.Text] = ...
 
-valuetype: google.protobuf.descriptor.FieldDescriptor = ...
+valuetype: google.protobuf.internal.extension_dict._ExtensionFieldDescriptor[google.protobuf.descriptor_pb2.FieldOptions, typing.Text] = ...

mypy_protobuf/main.py

@@ -376,18 +376,7 @@ def write_messages(
                         )
                     l("")
 
-                for ext in desc.extension:
-                    l(
-                        "{}: {}[{}, {}] = ...",
-                        ext.name,
-                        self._import(
-                            "google.protobuf.internal.extension_dict",
-                            "_ExtensionFieldDescriptor",
-                        ),
-                        self._import_message(ext.extendee),
-                        self.python_type(ext),
-                    )
-                    l("")
+                self.write_extensions(desc.extension)
 
                 # Constructor
                 self_arg = "self_" if any(f.name == "self" for f in fields) else "self"
@@ -522,14 +511,18 @@ def write_stringly_typed_fields(self, desc: d.DescriptorProto) -> None:
             )
 
     def write_extensions(self, extensions: Sequence[d.FieldDescriptorProto]) -> None:
-        if not extensions:
-            return
         l = self._write_line
-        field_descriptor_class = self._import(
-            "google.protobuf.descriptor", "FieldDescriptor"
-        )
-        for extension in extensions:
-            l("{}: {} = ...", extension.name, field_descriptor_class)
+        for ext in extensions:
+            l(
+                "{}: {}[{}, {}] = ...",
+                ext.name,
+                self._import(
+                    "google.protobuf.internal.extension_dict",
+                    "_ExtensionFieldDescriptor",
+                ),
+                self._import_message(ext.extendee),
+                self.python_type(ext),
+            )
             l("")
 
     def write_methods(
@@ -706,7 +699,7 @@ def python_type(self, field: d.FieldDescriptorProto) -> str:
         if casttype:
             return self._import_casttype(casttype)
 
-        mapping: Dict[d.FieldDescriptorProto.TypeValue, Callable[[], str]] = {
+        mapping: Dict[d.FieldDescriptorProto.Type.V, Callable[[], str]] = {
             d.FieldDescriptorProto.TYPE_DOUBLE: lambda: self._builtin("float"),
             d.FieldDescriptorProto.TYPE_FLOAT: lambda: self._builtin("float"),
             d.FieldDescriptorProto.TYPE_INT64: lambda: self._builtin("int"),

proto/testproto/test_extensions3.proto

@@ -11,3 +11,7 @@ extend google.protobuf.FieldOptions {
 extend google.protobuf.MessageOptions {
   string test_message_option = 51234;
 }
+
+message MessageOptionsTestMsg {
+    option (test_message_option) = "Hello world!";
+}

requirements.txt

@@ -15,4 +15,4 @@ grpcio-tools==1.38.1 ; python_version>="3.0"
 grpc-stubs>=1.24.6 ; python_version>="3.0"
 
 types-six==0.1.7
-types-protobuf==0.1.14
+types-protobuf==3.17.1

run_test.sh

@@ -95,6 +95,9 @@ find test/generated -type f -not \( -name "*.expected" -or -name "__init__.py" \
     else
         FILES=$FILES38
     fi
+    if [ -e $CUSTOM_TYPESHED_DIR ]; then
+        export MYPYPATH=$CUSTOM_TYPESHED_DIR/stubs/protobuf
+    fi
     mypy --strict --custom-typeshed-dir=$CUSTOM_TYPESHED_DIR --python-version=$PY_VER_MYPY_TARGET --pretty --show-error-codes $FILES
 
     # run mypy on negative-tests (expected mypy failures)

test/generated/mypy_protobuf/extensions_pb2.pyi.expected

@@ -3,11 +3,14 @@
 isort:skip_file
 """
 import google.protobuf.descriptor
+import google.protobuf.descriptor_pb2
+import google.protobuf.internal.extension_dict
+import typing
 
 DESCRIPTOR: google.protobuf.descriptor.FileDescriptor = ...
 
-casttype: google.protobuf.descriptor.FieldDescriptor = ...
+casttype: google.protobuf.internal.extension_dict._ExtensionFieldDescriptor[google.protobuf.descriptor_pb2.FieldOptions, typing.Text] = ...
 
-keytype: google.protobuf.descriptor.FieldDescriptor = ...
+keytype: google.protobuf.internal.extension_dict._ExtensionFieldDescriptor[google.protobuf.descriptor_pb2.FieldOptions, typing.Text] = ...
 
-valuetype: google.protobuf.descriptor.FieldDescriptor = ...
+valuetype: google.protobuf.internal.extension_dict._ExtensionFieldDescriptor[google.protobuf.descriptor_pb2.FieldOptions, typing.Text] = ...

test/generated/testproto/test_extensions3_pb2.pyi.expected

@@ -3,9 +3,20 @@
 isort:skip_file
 """
 import google.protobuf.descriptor
+import google.protobuf.descriptor_pb2
+import google.protobuf.internal.extension_dict
+import google.protobuf.message
+import typing
 
 DESCRIPTOR: google.protobuf.descriptor.FileDescriptor = ...
 
-test_field_extension: google.protobuf.descriptor.FieldDescriptor = ...
+class MessageOptionsTestMsg(google.protobuf.message.Message):
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor = ...
 
-test_message_option: google.protobuf.descriptor.FieldDescriptor = ...
+    def __init__(self,
+        ) -> None: ...
+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] = ...

test/test_generated_mypy.py

@@ -40,6 +40,7 @@
     OuterMessage3,
     SimpleProto3,
 )
+from testproto.test_extensions3_pb2 import MessageOptionsTestMsg, test_message_option
 from testproto.Capitalized.Capitalized_pb2 import lower, lower2, Upper
 
 from typing import (
@@ -149,8 +150,8 @@ def grab_expectations(filename, marker):
     assert errors_38 == expected_errors_38
 
     # Some sanity checks to make sure we don't mess this up. Please update as necessary.
-    assert len(errors_27) == 53
-    assert len(errors_38) == 65
+    assert len(errors_27) == 58
+    assert len(errors_38) == 70
 
 
 def test_func():
@@ -438,6 +439,14 @@ def test_extensions_proto2():
     assert len(s2.Extensions) == 1
 
 
+def test_extensions_proto3():
+    # type: () -> None
+    assert (
+        MessageOptionsTestMsg.DESCRIPTOR.GetOptions().Extensions[test_message_option]
+        == "Hello world!"
+    )
+
+
 def test_constructor_proto2():
     # type: () -> None
     x = Simple2()  # It's OK to omit a required field from the constructor.

test_negative/negative.py

@@ -94,14 +94,11 @@
 # changes to typeshed a few months after the 1.24 release
 # See https://github.com/python/typeshed/pull/4833
 _ = s.Extensions["foo"]  # E:2.7 E:3.8
-# TODO - these will give errors once again once we are able to
-# revert https://github.com/python/typeshed/pull/4833 later on
-# a few months after 1.24 releases
-_ = s.Extensions[SeparateFileExtension.ext]
-_ = SeparateFileExtension.ext in s.Extensions
-del s.Extensions[SeparateFileExtension.ext]
-s.HasExtension(SeparateFileExtension.ext)
-simple2.ClearExtension(Extensions1.ext)
+_ = s.Extensions[SeparateFileExtension.ext]  # E:2.7 E:3.8
+_ = SeparateFileExtension.ext in s.Extensions  # E:2.7 E:3.8
+del s.Extensions[SeparateFileExtension.ext]  # E:2.7 E:3.8
+s.HasExtension(SeparateFileExtension.ext)  # E:2.7 E:3.8
+simple2.ClearExtension(Extensions1.ext)  # E:2.7 E:3.8
 
 
 for x in s.Extensions:

test_negative/output.expected.2.7

@@ -27,45 +27,47 @@ test_negative/negative.py:88: error: Incompatible types in assignment (expressio
 test_negative/negative.py:89: error: "Type[Extensions1]" has no attribute "bad"
 test_negative/negative.py:91: error: "Extensions1" has no attribute "foo"
 test_negative/negative.py:92: error: Incompatible types in assignment (expression has type "Extensions2", variable has type "Extensions1")
-test_negative/negative.py:96: error: No overload variant of "__getitem__" of "_ExtensionDict" matches argument type "str"
-test_negative/negative.py:96: note: Possible overload variants:
-test_negative/negative.py:96: note:     def [_ExtenderMessageT <: Message] __getitem__(self, _ExtensionFieldDescriptor[Simple1, _ExtenderMessageT]) -> _ExtenderMessageT
-test_negative/negative.py:96: note:     def __getitem__(self, FieldDescriptor) -> Any
-test_negative/negative.py:109: error: Incompatible types in assignment (expression has type "int", variable has type "_ExtensionFieldDescriptor[Simple1, Any]")
-test_negative/negative.py:113: error: Incompatible types in assignment (expression has type "Union[Literal[u'b_oneof_1'], Literal[u'b_oneof_2'], None]", variable has type "Union[Literal[u'a_oneof_1'], Literal[u'a_oneof_2'], Literal[u'outer_message_in_oneof'], Literal[u'outer_enum_in_oneof'], Literal[u'inner_enum_in_oneof'], None]")
-test_negative/negative.py:116: error: "Descriptor" has no attribute "Garbage"
-test_negative/negative.py:117: error: "Descriptor" has no attribute "Garbage"
-test_negative/negative.py:120: error: "EnumDescriptor" has no attribute "Garbage"
-test_negative/negative.py:123: error: "FileDescriptor" has no attribute "Garbage"
-test_negative/negative.py:130: error: "V" has no attribute "FOO"
-test_negative/negative.py:134: error: Argument 1 to "Name" of "_EnumTypeWrapper" has incompatible type "int"; expected "V"
-test_negative/negative.py:136: error: Argument 1 to "Name" of "_EnumTypeWrapper" has incompatible type "testproto.test3_pb2.SimpleProto3.InnerEnum.V"; expected "testproto.test3_pb2.OuterEnum.V"
-test_negative/negative.py:138: error: Argument 1 to "Name" of "_EnumTypeWrapper" has incompatible type "testproto.test3_pb2.SimpleProto3.InnerEnum.V"; expected "testproto.test3_pb2.OuterEnum.V"
-test_negative/negative.py:140: error: Argument 1 to "Name" of "_EnumTypeWrapper" has incompatible type "testproto.test3_pb2.SimpleProto3.InnerEnum.V"; expected "testproto.test3_pb2.OuterEnum.V"
-test_negative/negative.py:144: error: "ScalarMap[int, unicode]" has no attribute "get_or_create"
-test_negative/negative.py:146: error: No overload variant of "get" of "Mapping" matches argument type "str"
-test_negative/negative.py:146: note: Possible overload variant:
-test_negative/negative.py:146: note:     def get(self, k: int) -> Optional[unicode]
-test_negative/negative.py:146: note:     <1 more non-matching overload not shown>
-test_negative/negative.py:147: error: No overload variant of "get" of "Mapping" matches argument type "str"
-test_negative/negative.py:147: note: Possible overload variant:
-test_negative/negative.py:147: note:     def get(self, k: int) -> Optional[OuterMessage3]
-test_negative/negative.py:147: note:     <1 more non-matching overload not shown>
-test_negative/negative.py:150: error: Incompatible types in assignment (expression has type "Optional[unicode]", variable has type "int")
-test_negative/negative.py:151: error: Incompatible types in assignment (expression has type "Optional[OuterMessage3]", variable has type "int")
-test_negative/negative.py:153: error: Dict entry 0 has incompatible type "str": "int"; expected "int": "unicode"
-test_negative/negative.py:153: error: Dict entry 0 has incompatible type "str": "str"; expected "int": "OuterMessage3"
-test_negative/negative.py:157: error: Incompatible types in assignment (expression has type "int", variable has type "UserId")
-test_negative/negative.py:158: error: Incompatible types in assignment (expression has type "str", variable has type "Email")
-test_negative/negative.py:159: error: Invalid index type "int" for "ScalarMap[UserId, Email]"; expected type "UserId"
-test_negative/negative.py:159: error: Incompatible types in assignment (expression has type "str", target has type "Email")
-test_negative/negative.py:160: error: Incompatible types in assignment (expression has type "str", target has type "Email")
-test_negative/negative.py:161: error: Invalid index type "int" for "ScalarMap[UserId, Email]"; expected type "UserId"
-test_negative/negative.py:162: error: Argument "user_id" to "Simple1" has incompatible type "int"; expected "Optional[UserId]"
-test_negative/negative.py:162: error: Argument "email" to "Simple1" has incompatible type "str"; expected "Optional[Email]"
-test_negative/negative.py:162: error: Dict entry 0 has incompatible type "int": "str"; expected "UserId": "Email"
-test_negative/negative.py:165: error: Module "testproto.reexport_pb2" has no attribute "Inner"
-test_negative/negative.py:169: error: Argument "a_bool" to "OuterMessage3" has incompatible type "None"; expected "bool"
-test_negative/negative.py:174: error: Property "a_repeated_string" defined in "Simple1" is read-only
-test_negative/negative.py:175: error: Property "rep_inner_enum" defined in "Simple1" is read-only
-Found 57 errors in 1 file (checked 17 source files)
+test_negative/negative.py:96: error: Invalid index type "str" for "_ExtensionDict[Simple1]"; expected type "_ExtensionFieldDescriptor[Simple1, <nothing>]"
+test_negative/negative.py:97: error: Invalid index type "_ExtensionFieldDescriptor[Simple2, SeparateFileExtension]" for "_ExtensionDict[Simple1]"; expected type "_ExtensionFieldDescriptor[Simple1, SeparateFileExtension]"
+test_negative/negative.py:98: error: Unsupported operand types for in ("_ExtensionFieldDescriptor[Simple2, SeparateFileExtension]" and "_ExtensionDict[Simple1]")
+test_negative/negative.py:99: error: Argument 1 to "__delitem__" of "_ExtensionDict" has incompatible type "_ExtensionFieldDescriptor[Simple2, SeparateFileExtension]"; expected "_ExtensionFieldDescriptor[Simple1, SeparateFileExtension]"
+test_negative/negative.py:100: error: Argument 1 to "HasExtension" of "Message" has incompatible type "_ExtensionFieldDescriptor[Simple2, SeparateFileExtension]"; expected "_ExtensionFieldDescriptor[Simple1, Any]"
+test_negative/negative.py:101: error: Argument 1 to "ClearExtension" of "Message" has incompatible type "_ExtensionFieldDescriptor[Simple1, Extensions1]"; expected "_ExtensionFieldDescriptor[Simple2, Any]"
+test_negative/negative.py:106: error: Incompatible types in assignment (expression has type "int", variable has type "_ExtensionFieldDescriptor[Simple1, Any]")
+test_negative/negative.py:110: error: Incompatible types in assignment (expression has type "Union[Literal[u'b_oneof_1'], Literal[u'b_oneof_2'], None]", variable has type "Union[Literal[u'a_oneof_1'], Literal[u'a_oneof_2'], Literal[u'outer_message_in_oneof'], Literal[u'outer_enum_in_oneof'], Literal[u'inner_enum_in_oneof'], None]")
+test_negative/negative.py:113: error: "Descriptor" has no attribute "Garbage"
+test_negative/negative.py:114: error: "Descriptor" has no attribute "Garbage"
+test_negative/negative.py:117: error: "EnumDescriptor" has no attribute "Garbage"
+test_negative/negative.py:120: error: "FileDescriptor" has no attribute "Garbage"
+test_negative/negative.py:127: error: "V" has no attribute "FOO"
+test_negative/negative.py:131: error: Argument 1 to "Name" of "_EnumTypeWrapper" has incompatible type "int"; expected "V"
+test_negative/negative.py:133: error: Argument 1 to "Name" of "_EnumTypeWrapper" has incompatible type "testproto.test3_pb2.SimpleProto3.InnerEnum.V"; expected "testproto.test3_pb2.OuterEnum.V"
+test_negative/negative.py:135: error: Argument 1 to "Name" of "_EnumTypeWrapper" has incompatible type "testproto.test3_pb2.SimpleProto3.InnerEnum.V"; expected "testproto.test3_pb2.OuterEnum.V"
+test_negative/negative.py:137: error: Argument 1 to "Name" of "_EnumTypeWrapper" has incompatible type "testproto.test3_pb2.SimpleProto3.InnerEnum.V"; expected "testproto.test3_pb2.OuterEnum.V"
+test_negative/negative.py:141: error: "ScalarMap[int, unicode]" has no attribute "get_or_create"
+test_negative/negative.py:143: error: No overload variant of "get" of "Mapping" matches argument type "str"
+test_negative/negative.py:143: note: Possible overload variant:
+test_negative/negative.py:143: note:     def get(self, k: int) -> Optional[unicode]
+test_negative/negative.py:143: note:     <1 more non-matching overload not shown>
+test_negative/negative.py:144: error: No overload variant of "get" of "Mapping" matches argument type "str"
+test_negative/negative.py:144: note: Possible overload variant:
+test_negative/negative.py:144: note:     def get(self, k: int) -> Optional[OuterMessage3]
+test_negative/negative.py:144: note:     <1 more non-matching overload not shown>
+test_negative/negative.py:147: error: Incompatible types in assignment (expression has type "Optional[unicode]", variable has type "int")
+test_negative/negative.py:148: error: Incompatible types in assignment (expression has type "Optional[OuterMessage3]", variable has type "int")
+test_negative/negative.py:150: error: Dict entry 0 has incompatible type "str": "int"; expected "int": "unicode"
+test_negative/negative.py:150: error: Dict entry 0 has incompatible type "str": "str"; expected "int": "OuterMessage3"
+test_negative/negative.py:154: error: Incompatible types in assignment (expression has type "int", variable has type "UserId")
+test_negative/negative.py:155: error: Incompatible types in assignment (expression has type "str", variable has type "Email")
+test_negative/negative.py:156: error: Invalid index type "int" for "ScalarMap[UserId, Email]"; expected type "UserId"
+test_negative/negative.py:156: error: Incompatible types in assignment (expression has type "str", target has type "Email")
+test_negative/negative.py:157: error: Incompatible types in assignment (expression has type "str", target has type "Email")
+test_negative/negative.py:158: error: Invalid index type "int" for "ScalarMap[UserId, Email]"; expected type "UserId"
+test_negative/negative.py:159: error: Argument "user_id" to "Simple1" has incompatible type "int"; expected "Optional[UserId]"
+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: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)