add aio support

[sagewe]

see #216

This pr add flag indicate mypy-protobuf should generate grpc stubs with aio.Server support

script like this

protoc \
    --python_out=output/location \
    --mypy_out=readable_stubs:output/location \
    --grpc_out=output/location \
    --mypy_grpc_out=aio,readable_stubs:output/location

generate from dummy.proto to

"""
@generated by mypy-protobuf.  Do not edit manually!
isort:skip_file
"""
from abc import (
    ABCMeta,
    abstractmethod,
)

from grpc import (
    Channel,
    Server,
    ServicerContext,
    StreamStreamMultiCallable,
    StreamUnaryMultiCallable,
    UnaryStreamMultiCallable,
    UnaryUnaryMultiCallable,
)

from grpc.aio import (
    Server as AioServer,
)

from typing import (
    Iterator,
    Union,
)


from .dummy_pb2 import *
class DummyServiceStub:
    def __init__(self, channel: Channel) -> None: ...
    UnaryUnary:UnaryUnaryMultiCallable[
        DummyRequest,
        DummyReply] = ...

    UnaryStream:UnaryStreamMultiCallable[
        DummyRequest,
        DummyReply] = ...

    StreamUnary:StreamUnaryMultiCallable[
        DummyRequest,
        DummyReply] = ...

    StreamStream:StreamStreamMultiCallable[
        DummyRequest,
        DummyReply] = ...


class DummyServiceServicer(metaclass=ABCMeta):
    @abstractmethod
    def UnaryUnary(self,
        request: DummyRequest,
        context: ServicerContext,
    ) -> DummyReply: ...

    @abstractmethod
    def UnaryStream(self,
        request: DummyRequest,
        context: ServicerContext,
    ) -> Iterator[DummyReply]: ...

    @abstractmethod
    def StreamUnary(self,
        request: Iterator[DummyRequest],
        context: ServicerContext,
    ) -> DummyReply: ...

    @abstractmethod
    def StreamStream(self,
        request: Iterator[DummyRequest],
        context: ServicerContext,
    ) -> Iterator[DummyReply]: ...


def add_DummyServiceServicer_to_server(servicer: DummyServiceServicer, server: Union[Server, AioServer]) -> None: ...

[nipunn1313]

looks lovely! Thanks for the PR!

From my reading of the docs of GRPC - it seems that AIO is trying to replace the original server

Seems like there are a few options here. I haven't directly used this GRPC generation, so will defer to folks on what makes the most sense here:

1. Generate the union unconditionally

def add_DummyServiceServicer_to_server(servicer: DummyServiceServicer, server: Union[Server, AioServer]) -> None: ...

2. Generate either Server or AioServer depending on the cmd line flag

def add_DummyServiceServicer_to_server(servicer: DummyServiceServicer, server: Server) -> None: ...
# w aio flag
def add_DummyServiceServicer_to_server(servicer: DummyServiceServicer, server: AioServer) -> None: ...

3. Generate Server by default - and union w/ the flag

def add_DummyServiceServicer_to_server(servicer: DummyServiceServicer, server: Server) -> None: ...
# w aio flag
def add_DummyServiceServicer_to_server(servicer: DummyServiceServicer, server: Union[Server, AioServer]) -> None: ...

As implemented - I believe it's option 3. I think I'd prefer either option 1 or option 2 (for simplicity / stronger typing). Option 3 appears to add complexity and weakens typing for the aio case. Are there cases in which we'd want to mix and match server types to servicers?

I'd be enthusiastic about starting with (1) and adding complexity as needed. More flags -> more complexity to maintain.
Additional benefit - is that the default flags are well tested.

Want to acknowledge that I don't understand the use case as well as you - so please poke holes in my reasoning here if you have another viewpoint. Thank you!

Action Items

• Use approach 1 (no flag -> always use union)
• Run tests - so that the ".expected" files update
• Add an entry to Changelog under "upcoming"
• Add your name to contributors (if you'd like)!

Thank you!
If this seems like a lot of work and you aren't interested / are too busy, let me know and I can put it on my backburner to do in the future.

[MHDante]

I would agree for using option 1, as this best represents the actual typing of the .py code. Furthermore, (as is my use case) I generate the files once and use them as a shared library across modules, some of which will use AIO and others will use Threading.

Also, this implementation seems to be incomplete, as using asyncIO should allow you to change the return value of the generated functions to a union of ReturnType and Awaitable[ReturnType]

It's important to keep the union in this case since the AIO package allows you to define the function in either format.

I will make some of these changes and update the PR.

[jack-tutor]

Has there been any movement on this; it would be very useful for me. I also want to second this point: Also, this implementation seems to be incomplete, as using asyncIO should allow you to change the return value of the generated functions to a union of ReturnType and Awaitable[ReturnType]

[MHDante]

Sorry, I've stepped away from this for a bit. At the moment my team is adding #type: ignore and cast where needed. I'll see if I can put some elbow grease into this on the weekend.

[nipunn1313]

I believe it depends on shabbyrobe/grpc-stubs#22 - so we can import from grpc.aio

I quickly tried to pick this up and got

test/generated/testproto/grpc/import_pb2_grpc.pyi:8: error: Skipping analyzing "grpc.aio": found module but no type hints or library stubs  [import]

I haven't been focusing on it. Feel free to pick up this diff and proceed. You'll probably have to co-develop in both shabbyrobe/grpc-stubs and here. I'm sure @shabbyrobe would be happy to review that side of it if you manage to get it all working.

[nipunn1313]

I believe we want to generate something like this

class DummyServiceServicer:
    def UnaryUnary(self,
        request: DummyRequest,
        context: grpc.ServicerContext,
    ) -> typing.Union[DummyReply, typing.Awaitable[DummyReply]]: ...

def add_DummyServiceServicer_to_server(
    servicer: DummyServiceServicer,
    server: typing.Union[grpc.Server, grpc.aio.AioServer]) -> None: ...

Even better would be something like this

A = TypeVar('A', typing.Awaitable, typing.Union)  # Using Union as a noop wrapper

class DummyServiceServicer(Generic[A]):
    def UnaryUnary(self,
        request: DummyRequest,
        context: grpc.ServicerContext,
    ) -> typing.Union[DummyReply, A[DummyReply]]: ...

@overload
def add_DummyServiceServicer_to_server(
    servicer: DummyServiceServicer[Awaitable],
    server: grpc.aio.AioServer) -> None: ...

@overload
def add_DummyServiceServicer_to_server(
    servicer: DummyServiceServicer[Union],
    server: grpc.Server) -> None: ...

Feels weird using Union as a noop wrapper in this way, but I couldn't think of another way.

[MHDante]

Ah, I remember now where I got stuck.

The aio package itself contains typings! 🎉

• however, some of the typings are incorrect
• some of the types depend on the non-aio types (there is some complexity in handling interop with grpc-stubs)

There seems to be some Issues that are blocking this:
grpc/grpc#26499 ❌
grpc/grpc#26500 ✔️
grpc/grpc#26498 ⌛

My intention was to do as much of the work as I could on the upstream repos as I could before hacking up shims here. This may not be necessary, and it would be useful to have some placeholder functionality until such time.

After some deliberation, It seemed impractical to combine the definition for AIO/non-AIO servicers. This is primarily because the ServicerContext that is passed into the functions has some different members. This becomes very obvious if you try to type server-streaming method with cancellation support. 😖 Making the parameter type a union would become impractical as service implementers would have to do type assertions/checks in every method.

I personally am using the same mypy-protobuf generated files in more than one project, some of which are now slowly transitioning to aio (some never will). I am trying to get a feel for all the nuances of the aio package at the moment, but they are less obvious than they seem. It would be advantageous to me to be able to use both types of services in a single set of typings.

I understand that this is not everybody's use case, and most people would be well served by a simple flag that compiles one or the other 😅 .

[nipunn1313]

After some deliberation, It seemed impractical to combine the definition for AIO/non-AIO servicers. This is primarily because the ServicerContext that is passed into the functions has some different members. This becomes very obvious if you try to type server-streaming method with cancellation support. 😖 Making the parameter type a union would become impractical as service implementers would have to do type assertions/checks in every method.

This makes sense.

One idea could be to add a mypy_protobuf proto extension as a service option.
Based on the content of the ServiceOption, the individual service could be created as Sync or Async. You could even theoretically do it on a method-by-method basis if that makes more sense.

Eg extend google.protobuf.ServiceOptions or extend google.protobuf.MethodOptions.

Then you'd declare your proto as something like

service DummyService {
  option (mypy_protobuf.async_service) = true;
  rpc UnaryUnary (DummyRequest) returns (DummyReply) {}
}

or

service DummyService {
  rpc UnaryUnary (DummyRequest) returns (DummyReply) {
      option (mypy_protobuf.async_service) = true;
  } 
}

https://developers.google.com/protocol-buffers/docs/proto#customoptions

There are several examples of this in pb-jelly (another project I worked on) https://github.com/dropbox/pb-jelly/blob/main/pb-jelly-gen/codegen/rust/extensions.proto - which codegens rust.

[nipunn1313]

I wrote up #276 - to give you a starting point for the ideas above. It's not really useful on its own, but just wanted to give you something so you wouldn't get stuck with the nuances of getting proto extensions working.

[MHDante]

Please excuse my flip-flopping. But I don't think that setting whether the service is sync/async in the protos is where we should do it? I would argue it should be a flag as @weiwee originally suggested. This is because the server and client (and indeed multiple different versions of each) can each have different implementations of the proto file and that should be supported. ( I would argue this is one of the strengths of grpc/proto 🦾 ).

If It's ok with you I'd like to suggest continuing with the compiler flag approach.

[nipunn1313]

That logic makes sense. The sync/async is a question of implementation, not of spec.

There's still a couple of shortcomings to the flag approach

• Flags are hard to maintain - can blow up test matrix.
• Flag must be issued on a per-invocation basis rather than on a per-service basis. Meaning you can't have multiple services within a file with different implementation types. There is at least a workaround in this case (split into two files and make multiple invocations to protoc)

Using flags seems like the lesser of evils here.

That being said - if during implementation you can find a way to do it without the flag, and somehow support both sync/async cases (perhaps via @overload and type generics) - that would be ideal.

[shonfeder]

Looks like we got very close here! Too bad it wasn't ready for me :D. If I can help move this forward, and it's not a huge lift, I'd be happy to do so.