diff --git a/mypy_protobuf/extensions_pb2.py b/mypy_protobuf/extensions_pb2.py index 745bc028..ca9cdddb 100644 --- a/mypy_protobuf/extensions_pb2.py +++ b/mypy_protobuf/extensions_pb2.py @@ -20,7 +20,7 @@ syntax='proto2', serialized_options=None, create_key=_descriptor._internal_create_key, - serialized_pb=b'\n\x1emypy_protobuf/extensions.proto\x12\rmypy_protobuf\x1a google/protobuf/descriptor.proto:1\n\x08\x63\x61sttype\x12\x1d.google.protobuf.FieldOptions\x18\xe0\xd4\x03 \x01(\t:0\n\x07keytype\x12\x1d.google.protobuf.FieldOptions\x18\xe2\xd4\x03 \x01(\t:2\n\tvaluetype\x12\x1d.google.protobuf.FieldOptions\x18\xe3\xd4\x03 \x01(\t' + serialized_pb=b'\n\x1emypy_protobuf/extensions.proto\x12\rmypy_protobuf\x1a google/protobuf/descriptor.proto:1\n\x08\x63\x61sttype\x12\x1d.google.protobuf.FieldOptions\x18\xe0\xd4\x03 \x01(\t:0\n\x07keytype\x12\x1d.google.protobuf.FieldOptions\x18\xe2\xd4\x03 \x01(\t:2\n\tvaluetype\x12\x1d.google.protobuf.FieldOptions\x18\xe3\xd4\x03 \x01(\t:8\n\rasync_service\x12\x1f.google.protobuf.ServiceOptions\x18\xe4\xd4\x03 \x01(\x08' , dependencies=[google_dot_protobuf_dot_descriptor__pb2.DESCRIPTOR,]) @@ -49,14 +49,24 @@ message_type=None, enum_type=None, containing_type=None, is_extension=True, extension_scope=None, serialized_options=None, file=DESCRIPTOR, create_key=_descriptor._internal_create_key) +ASYNC_SERVICE_FIELD_NUMBER = 60004 +async_service = _descriptor.FieldDescriptor( + name='async_service', full_name='mypy_protobuf.async_service', index=3, + number=60004, type=8, cpp_type=7, label=1, + has_default_value=False, default_value=False, + message_type=None, enum_type=None, containing_type=None, + is_extension=True, extension_scope=None, + serialized_options=None, file=DESCRIPTOR, create_key=_descriptor._internal_create_key) DESCRIPTOR.extensions_by_name['casttype'] = casttype DESCRIPTOR.extensions_by_name['keytype'] = keytype DESCRIPTOR.extensions_by_name['valuetype'] = valuetype +DESCRIPTOR.extensions_by_name['async_service'] = async_service _sym_db.RegisterFileDescriptor(DESCRIPTOR) google_dot_protobuf_dot_descriptor__pb2.FieldOptions.RegisterExtension(casttype) google_dot_protobuf_dot_descriptor__pb2.FieldOptions.RegisterExtension(keytype) google_dot_protobuf_dot_descriptor__pb2.FieldOptions.RegisterExtension(valuetype) +google_dot_protobuf_dot_descriptor__pb2.ServiceOptions.RegisterExtension(async_service) # @@protoc_insertion_point(module_scope) diff --git a/mypy_protobuf/extensions_pb2.pyi b/mypy_protobuf/extensions_pb2.pyi index d59b931c..92f282a3 100644 --- a/mypy_protobuf/extensions_pb2.pyi +++ b/mypy_protobuf/extensions_pb2.pyi @@ -2,6 +2,7 @@ @generated by mypy-protobuf. Do not edit manually! isort:skip_file """ +import builtins import google.protobuf.descriptor import google.protobuf.descriptor_pb2 import google.protobuf.internal.extension_dict @@ -17,3 +18,6 @@ keytype: google.protobuf.internal.extension_dict._ExtensionFieldDescriptor[googl # Tells mypy to use a specific type for values; only makes sense on map fields valuetype: google.protobuf.internal.extension_dict._ExtensionFieldDescriptor[google.protobuf.descriptor_pb2.FieldOptions, typing.Text] = ... + +# Tells mypy to generate async methods for this service +async_service: google.protobuf.internal.extension_dict._ExtensionFieldDescriptor[google.protobuf.descriptor_pb2.ServiceOptions, builtins.bool] = ... diff --git a/mypy_protobuf/main.py b/mypy_protobuf/main.py index bc662cbb..4d5e7c91 100644 --- a/mypy_protobuf/main.py +++ b/mypy_protobuf/main.py @@ -681,9 +681,14 @@ def _input_type( return result def _output_type( - self, method: d.MethodDescriptorProto, use_stream_iterator: bool = True + self, + method: d.MethodDescriptorProto, + async_service: bool, + use_stream_iterator: bool = True, ) -> str: result = self._import_message(method.output_type) + if async_service: + result = "{}[{}]".format(self._import("typing", "Awaitable"), result) if use_stream_iterator and method.server_streaming: result = "{}[{}]".format(self._import("typing", "Iterator"), result) return result @@ -700,6 +705,7 @@ def write_grpc_methods( if not methods: l("pass") l("") + async_service = service.options.Extensions[extensions_pb2.async_service] for i, method in methods: scl = scl_prefix + [d.ServiceDescriptorProto.METHOD_FIELD_NUMBER, i] self._write_comments(scl) @@ -709,7 +715,7 @@ def write_grpc_methods( with self._indent(): l("request: {},", self._input_type(method)) l("context: {},", self._import("grpc", "ServicerContext")) - l(") -> {}: ...", self._output_type(method)) + l(") -> {}: ...", self._output_type(method, async_service)) l("") def write_grpc_stub_methods( @@ -724,6 +730,7 @@ def write_grpc_stub_methods( if not methods: l("pass") l("") + async_service = service.options.Extensions[extensions_pb2.async_service] for i, method in methods: scl = scl_prefix + [d.ServiceDescriptorProto.METHOD_FIELD_NUMBER, i] self._write_comments(scl) @@ -731,7 +738,7 @@ def write_grpc_stub_methods( l("{}:{}[", method.name, self._callable_type(method)) with self._indent(): l("{},", self._input_type(method, False)) - l("{}] = ...", self._output_type(method, False)) + l("{}] = ...", self._output_type(method, async_service, False)) l("") def write_grpc_services( @@ -749,6 +756,8 @@ def write_grpc_services( if service.name in PYTHON_RESERVED: continue + async_service = service.options.Extensions[extensions_pb2.async_service] + scl = scl_prefix + [i] self._write_comments(scl) @@ -776,7 +785,9 @@ def write_grpc_services( "def add_{}Servicer_to_server(servicer: {}Servicer, server: {}) -> None: ...", service.name, service.name, - self._import("grpc", "Server"), + self._import("grpc", "Server") + if async_service + else self._import("grpc", "Server"), ) l("") diff --git a/proto/mypy_protobuf/extensions.proto b/proto/mypy_protobuf/extensions.proto index f0294451..df59bbb3 100644 --- a/proto/mypy_protobuf/extensions.proto +++ b/proto/mypy_protobuf/extensions.proto @@ -13,3 +13,8 @@ extend google.protobuf.FieldOptions { // Tells mypy to use a specific type for values; only makes sense on map fields optional string valuetype = 60003; } + +extend google.protobuf.ServiceOptions { + // Tells mypy to generate async methods for this service + optional bool async_service = 60004; +} diff --git a/proto/testproto/grpc/dummy.proto b/proto/testproto/grpc/dummy.proto index ed626809..e8c4597c 100644 --- a/proto/testproto/grpc/dummy.proto +++ b/proto/testproto/grpc/dummy.proto @@ -3,6 +3,8 @@ syntax = "proto3"; package dummy; +import "mypy_protobuf/extensions.proto"; + message DummyRequest { string value = 1; } @@ -16,4 +18,9 @@ service DummyService { rpc UnaryStream (DummyRequest) returns (stream DummyReply) {} rpc StreamUnary (stream DummyRequest) returns (DummyReply) {} rpc StreamStream (stream DummyRequest) returns (stream DummyReply) {} -} \ No newline at end of file +} + +service DummyAsyncService { + option (mypy_protobuf.async_service) = true; + rpc UnaryUnary (DummyRequest) returns (DummyReply) {} +} diff --git a/run_test.sh b/run_test.sh index e5d9205a..a05b06bd 100755 --- a/run_test.sh +++ b/run_test.sh @@ -75,6 +75,11 @@ MYPY_PROTOBUF_VENV=venv_$PY_VER_MYPY_PROTOBUF test "$(protoc-gen-mypy_grpc -V)" = "mypy-protobuf 2.9" test "$(protoc-gen-mypy_grpc --version)" = "mypy-protobuf 2.9" + # Compile protoc -> mypy using mypy_protobuf + # Prereq - create the mypy.proto python proto + $PROTOC $PROTOC_ARGS --python_out=. `find proto/mypy_protobuf -name "*.proto"` + $PROTOC $PROTOC_ARGS --mypy_out=. `find proto/mypy_protobuf -name "*.proto"` + # Run mypy on mypy-protobuf internal code for developers to catch issues FILES="mypy_protobuf/main.py" $MYPY_VENV/bin/mypy --custom-typeshed-dir=$CUSTOM_TYPESHED_DIR --python-executable=$MYPY_PROTOBUF_VENV/bin/python3 --python-version=$PY_VER_MYPY_PROTOBUF_SHORT $FILES @@ -96,11 +101,6 @@ MYPY_PROTOBUF_VENV=venv_$PY_VER_MYPY_PROTOBUF # Compile protoc -> python $PROTOC $PROTOC_ARGS --python_out=test/generated `find proto -name "*.proto"` - # Compile protoc -> mypy using mypy_protobuf - # Prereq - create the mypy.proto python proto - $PROTOC $PROTOC_ARGS --python_out=. `find proto/mypy_protobuf -name "*.proto"` - $PROTOC $PROTOC_ARGS --mypy_out=. `find proto/mypy_protobuf -name "*.proto"` - # Sanity check that our flags work $PROTOC $PROTOC_ARGS --mypy_out=quiet:test/generated `find proto -name "*.proto"` $PROTOC $PROTOC_ARGS --mypy_out=readable_stubs:test/generated `find proto -name "*.proto"` diff --git a/test/generated/mypy_protobuf/extensions_pb2.pyi b/test/generated/mypy_protobuf/extensions_pb2.pyi index d59b931c..92f282a3 100644 --- a/test/generated/mypy_protobuf/extensions_pb2.pyi +++ b/test/generated/mypy_protobuf/extensions_pb2.pyi @@ -2,6 +2,7 @@ @generated by mypy-protobuf. Do not edit manually! isort:skip_file """ +import builtins import google.protobuf.descriptor import google.protobuf.descriptor_pb2 import google.protobuf.internal.extension_dict @@ -17,3 +18,6 @@ keytype: google.protobuf.internal.extension_dict._ExtensionFieldDescriptor[googl # Tells mypy to use a specific type for values; only makes sense on map fields valuetype: google.protobuf.internal.extension_dict._ExtensionFieldDescriptor[google.protobuf.descriptor_pb2.FieldOptions, typing.Text] = ... + +# Tells mypy to generate async methods for this service +async_service: google.protobuf.internal.extension_dict._ExtensionFieldDescriptor[google.protobuf.descriptor_pb2.ServiceOptions, builtins.bool] = ... diff --git a/test/generated/testproto/grpc/dummy_pb2_grpc.pyi b/test/generated/testproto/grpc/dummy_pb2_grpc.pyi index 310ee53f..2fd3bf42 100644 --- a/test/generated/testproto/grpc/dummy_pb2_grpc.pyi +++ b/test/generated/testproto/grpc/dummy_pb2_grpc.pyi @@ -53,3 +53,20 @@ class DummyServiceServicer(metaclass=abc.ABCMeta): def add_DummyServiceServicer_to_server(servicer: DummyServiceServicer, server: grpc.Server) -> None: ... + +class DummyAsyncServiceStub: + def __init__(self, channel: grpc.Channel) -> None: ... + UnaryUnary:grpc.UnaryUnaryMultiCallable[ + global___DummyRequest, + typing.Awaitable[global___DummyReply]] = ... + + +class DummyAsyncServiceServicer(metaclass=abc.ABCMeta): + @abc.abstractmethod + def UnaryUnary(self, + request: global___DummyRequest, + context: grpc.ServicerContext, + ) -> typing.Awaitable[global___DummyReply]: ... + + +def add_DummyAsyncServiceServicer_to_server(servicer: DummyAsyncServiceServicer, server: grpc.Server) -> None: ...