Skip to content

Commit 718f414

Browse files
committed
Barebones grpc aio attempt
1 parent bd793af commit 718f414

File tree

8 files changed

+69
-11
lines changed

8 files changed

+69
-11
lines changed

mypy_protobuf/extensions_pb2.py

+11-1
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

mypy_protobuf/extensions_pb2.pyi

+4
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
@generated by mypy-protobuf. Do not edit manually!
33
isort:skip_file
44
"""
5+
import builtins
56
import google.protobuf.descriptor
67
import google.protobuf.descriptor_pb2
78
import google.protobuf.internal.extension_dict
@@ -17,3 +18,6 @@ keytype: google.protobuf.internal.extension_dict._ExtensionFieldDescriptor[googl
1718

1819
# Tells mypy to use a specific type for values; only makes sense on map fields
1920
valuetype: google.protobuf.internal.extension_dict._ExtensionFieldDescriptor[google.protobuf.descriptor_pb2.FieldOptions, typing.Text] = ...
21+
22+
# Tells mypy to generate async methods for this service
23+
async_service: google.protobuf.internal.extension_dict._ExtensionFieldDescriptor[google.protobuf.descriptor_pb2.ServiceOptions, builtins.bool] = ...

mypy_protobuf/main.py

+15-4
Original file line numberDiff line numberDiff line change
@@ -681,9 +681,14 @@ def _input_type(
681681
return result
682682

683683
def _output_type(
684-
self, method: d.MethodDescriptorProto, use_stream_iterator: bool = True
684+
self,
685+
method: d.MethodDescriptorProto,
686+
async_service: bool,
687+
use_stream_iterator: bool = True,
685688
) -> str:
686689
result = self._import_message(method.output_type)
690+
if async_service:
691+
result = "{}[{}]".format(self._import("typing", "Awaitable"), result)
687692
if use_stream_iterator and method.server_streaming:
688693
result = "{}[{}]".format(self._import("typing", "Iterator"), result)
689694
return result
@@ -700,6 +705,7 @@ def write_grpc_methods(
700705
if not methods:
701706
l("pass")
702707
l("")
708+
async_service = service.options.Extensions[extensions_pb2.async_service]
703709
for i, method in methods:
704710
scl = scl_prefix + [d.ServiceDescriptorProto.METHOD_FIELD_NUMBER, i]
705711
self._write_comments(scl)
@@ -709,7 +715,7 @@ def write_grpc_methods(
709715
with self._indent():
710716
l("request: {},", self._input_type(method))
711717
l("context: {},", self._import("grpc", "ServicerContext"))
712-
l(") -> {}: ...", self._output_type(method))
718+
l(") -> {}: ...", self._output_type(method, async_service))
713719
l("")
714720

715721
def write_grpc_stub_methods(
@@ -724,14 +730,15 @@ def write_grpc_stub_methods(
724730
if not methods:
725731
l("pass")
726732
l("")
733+
async_service = service.options.Extensions[extensions_pb2.async_service]
727734
for i, method in methods:
728735
scl = scl_prefix + [d.ServiceDescriptorProto.METHOD_FIELD_NUMBER, i]
729736
self._write_comments(scl)
730737

731738
l("{}:{}[", method.name, self._callable_type(method))
732739
with self._indent():
733740
l("{},", self._input_type(method, False))
734-
l("{}] = ...", self._output_type(method, False))
741+
l("{}] = ...", self._output_type(method, async_service, False))
735742
l("")
736743

737744
def write_grpc_services(
@@ -749,6 +756,8 @@ def write_grpc_services(
749756
if service.name in PYTHON_RESERVED:
750757
continue
751758

759+
async_service = service.options.Extensions[extensions_pb2.async_service]
760+
752761
scl = scl_prefix + [i]
753762
self._write_comments(scl)
754763

@@ -776,7 +785,9 @@ def write_grpc_services(
776785
"def add_{}Servicer_to_server(servicer: {}Servicer, server: {}) -> None: ...",
777786
service.name,
778787
service.name,
779-
self._import("grpc", "Server"),
788+
self._import("grpc", "Server")
789+
if async_service
790+
else self._import("grpc", "Server"),
780791
)
781792
l("")
782793

proto/mypy_protobuf/extensions.proto

+5
Original file line numberDiff line numberDiff line change
@@ -13,3 +13,8 @@ extend google.protobuf.FieldOptions {
1313
// Tells mypy to use a specific type for values; only makes sense on map fields
1414
optional string valuetype = 60003;
1515
}
16+
17+
extend google.protobuf.ServiceOptions {
18+
// Tells mypy to generate async methods for this service
19+
optional bool async_service = 60004;
20+
}

proto/testproto/grpc/dummy.proto

+8-1
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,8 @@ syntax = "proto3";
33

44
package dummy;
55

6+
import "mypy_protobuf/extensions.proto";
7+
68
message DummyRequest {
79
string value = 1;
810
}
@@ -16,4 +18,9 @@ service DummyService {
1618
rpc UnaryStream (DummyRequest) returns (stream DummyReply) {}
1719
rpc StreamUnary (stream DummyRequest) returns (DummyReply) {}
1820
rpc StreamStream (stream DummyRequest) returns (stream DummyReply) {}
19-
}
21+
}
22+
23+
service DummyAsyncService {
24+
option (mypy_protobuf.async_service) = true;
25+
rpc UnaryUnary (DummyRequest) returns (DummyReply) {}
26+
}

run_test.sh

+5-5
Original file line numberDiff line numberDiff line change
@@ -75,6 +75,11 @@ MYPY_PROTOBUF_VENV=venv_$PY_VER_MYPY_PROTOBUF
7575
test "$(protoc-gen-mypy_grpc -V)" = "mypy-protobuf 2.9"
7676
test "$(protoc-gen-mypy_grpc --version)" = "mypy-protobuf 2.9"
7777

78+
# Compile protoc -> mypy using mypy_protobuf
79+
# Prereq - create the mypy.proto python proto
80+
$PROTOC $PROTOC_ARGS --python_out=. `find proto/mypy_protobuf -name "*.proto"`
81+
$PROTOC $PROTOC_ARGS --mypy_out=. `find proto/mypy_protobuf -name "*.proto"`
82+
7883
# Run mypy on mypy-protobuf internal code for developers to catch issues
7984
FILES="mypy_protobuf/main.py"
8085
$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
96101
# Compile protoc -> python
97102
$PROTOC $PROTOC_ARGS --python_out=test/generated `find proto -name "*.proto"`
98103

99-
# Compile protoc -> mypy using mypy_protobuf
100-
# Prereq - create the mypy.proto python proto
101-
$PROTOC $PROTOC_ARGS --python_out=. `find proto/mypy_protobuf -name "*.proto"`
102-
$PROTOC $PROTOC_ARGS --mypy_out=. `find proto/mypy_protobuf -name "*.proto"`
103-
104104
# Sanity check that our flags work
105105
$PROTOC $PROTOC_ARGS --mypy_out=quiet:test/generated `find proto -name "*.proto"`
106106
$PROTOC $PROTOC_ARGS --mypy_out=readable_stubs:test/generated `find proto -name "*.proto"`

test/generated/mypy_protobuf/extensions_pb2.pyi

+4
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
@generated by mypy-protobuf. Do not edit manually!
33
isort:skip_file
44
"""
5+
import builtins
56
import google.protobuf.descriptor
67
import google.protobuf.descriptor_pb2
78
import google.protobuf.internal.extension_dict
@@ -17,3 +18,6 @@ keytype: google.protobuf.internal.extension_dict._ExtensionFieldDescriptor[googl
1718

1819
# Tells mypy to use a specific type for values; only makes sense on map fields
1920
valuetype: google.protobuf.internal.extension_dict._ExtensionFieldDescriptor[google.protobuf.descriptor_pb2.FieldOptions, typing.Text] = ...
21+
22+
# Tells mypy to generate async methods for this service
23+
async_service: google.protobuf.internal.extension_dict._ExtensionFieldDescriptor[google.protobuf.descriptor_pb2.ServiceOptions, builtins.bool] = ...

test/generated/testproto/grpc/dummy_pb2_grpc.pyi

+17
Original file line numberDiff line numberDiff line change
@@ -53,3 +53,20 @@ class DummyServiceServicer(metaclass=abc.ABCMeta):
5353

5454

5555
def add_DummyServiceServicer_to_server(servicer: DummyServiceServicer, server: grpc.Server) -> None: ...
56+
57+
class DummyAsyncServiceStub:
58+
def __init__(self, channel: grpc.Channel) -> None: ...
59+
UnaryUnary:grpc.UnaryUnaryMultiCallable[
60+
global___DummyRequest,
61+
typing.Awaitable[global___DummyReply]] = ...
62+
63+
64+
class DummyAsyncServiceServicer(metaclass=abc.ABCMeta):
65+
@abc.abstractmethod
66+
def UnaryUnary(self,
67+
request: global___DummyRequest,
68+
context: grpc.ServicerContext,
69+
) -> typing.Awaitable[global___DummyReply]: ...
70+
71+
72+
def add_DummyAsyncServiceServicer_to_server(servicer: DummyAsyncServiceServicer, server: grpc.Server) -> None: ...

0 commit comments

Comments
 (0)