Skip to content

Unimplemented plugin #593

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 7 commits into from
Apr 29, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions Makefile
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ typecheck:

_buf: clean
rm -rf src/viam/gen
chmod +x plugins/main.py
buf generate buf.build/viamrobotics/api
buf generate buf.build/viamrobotics/goutils
protol -e googl* --in-place -s _grpc.py -s _pb2.py -s _pb2.pyi -o src/viam/gen buf buf.build/viamrobotics/api
Expand Down
3 changes: 2 additions & 1 deletion buf.gen.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,8 @@ version: v1
plugins:
- name: python
out: ./src/viam/gen
- name: grpclib_python
- name: viam_grpc_python
out: ./src/viam/gen
path: [python3, ./plugin/main.py]
- name: mypy
out: ./src/viam/gen
2 changes: 2 additions & 0 deletions etc/generate_proto_import.py
Original file line number Diff line number Diff line change
Expand Up @@ -122,6 +122,8 @@ def check_class(obj) -> bool:
mod_name = f"viam.{PROTO_GEN_PACKAGE}.{package}.{imp}"
module = importlib.import_module(mod_name)
for name, _ in inspect.getmembers(module, check_class):
if name[0] == "_":
continue
class_names.append(name)

if class_names:
Expand Down
279 changes: 279 additions & 0 deletions plugin/main.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,279 @@
import os
import re
import sys

from typing import List, Any, Collection, Iterator, NamedTuple, cast
from typing import Dict, Tuple, Optional, Deque
from contextlib import contextmanager
from collections import deque

from google.protobuf.descriptor_pb2 import FileDescriptorProto, DescriptorProto
from google.protobuf.compiler.plugin_pb2 import CodeGeneratorRequest
from google.protobuf.compiler.plugin_pb2 import CodeGeneratorResponse

from grpclib import const
from grpclib import client
from grpclib import server
from grpclib import exceptions


_CARDINALITY = {
(False, False): const.Cardinality.UNARY_UNARY,
(True, False): const.Cardinality.STREAM_UNARY,
(False, True): const.Cardinality.UNARY_STREAM,
(True, True): const.Cardinality.STREAM_STREAM,
}


class Method(NamedTuple):
name: str
cardinality: const.Cardinality
request_type: str
reply_type: str


class Service(NamedTuple):
name: str
methods: List[Method]


class Buffer:
def __init__(self) -> None:
self._lines: List[str] = []
self._indent = 0

def add(self, string: str, *args: Any, **kwargs: Any) -> None:
line = " " * self._indent * 4 + string.format(*args, **kwargs)
self._lines.append(line.rstrip(" "))

@contextmanager
def indent(self) -> Iterator[None]:
self._indent += 1
try:
yield
finally:
self._indent -= 1

def content(self) -> str:
return "\n".join(self._lines) + "\n"


def render(
proto_file: str,
package: str,
imports: Collection[str],
services: Collection[Service],
) -> str:
buf = Buffer()
buf.add("# Generated by the Protocol Buffers compiler. DO NOT EDIT!")
buf.add("# source: {}", proto_file)
buf.add("# plugin: {}", __name__)
if not services:
return buf.content()

buf.add("import abc")
buf.add("import typing")
buf.add("")
buf.add("import {}", const.__name__)
buf.add("import {}", client.__name__)
buf.add("import {}", exceptions.__name__)
buf.add("if typing.TYPE_CHECKING:")
with buf.indent():
buf.add("import {}", server.__name__)

buf.add("")
for mod in imports:
buf.add("import {}", mod)
for service in services:
if package:
service_name = "{}.{}".format(package, service.name)
else:
service_name = service.name
buf.add("")
buf.add("")
buf.add("class {}Base(abc.ABC):", service.name)
with buf.indent():
for name, _, request_type, reply_type in service.methods:
buf.add("")
buf.add("@abc.abstractmethod")
buf.add(
"async def {}(self, stream: '{}.{}[{}, {}]') -> None:",
name,
server.__name__,
server.Stream.__name__,
request_type,
reply_type,
)
with buf.indent():
buf.add("pass")
buf.add("")
buf.add("def __mapping__(self) -> typing.Dict[str, {}.{}]:", const.__name__, const.Handler.__name__)
with buf.indent():
buf.add("return {{")
with buf.indent():
for method in service.methods:
name, cardinality, request_type, reply_type = method
full_name = "/{}/{}".format(service_name, name)
buf.add("'{}': {}.{}(", full_name, const.__name__, const.Handler.__name__)
with buf.indent():
buf.add("self.{},", name)
buf.add("{}.{}.{},", const.__name__, const.Cardinality.__name__, cardinality.name)
buf.add("{},", request_type)
buf.add("{},", reply_type)
buf.add("),")
buf.add("}}")

buf.add("")
buf.add("")
buf.add("class Unimplemented{}Base({}Base):", service.name, service.name)
with buf.indent():
for name, _, request_type, reply_type in service.methods:
buf.add("")
buf.add(
"async def {}(self, stream: '{}.{}[{}, {}]') -> None:",
name,
server.__name__,
server.Stream.__name__,
request_type,
reply_type,
)
with buf.indent():
buf.add("raise grpclib.exceptions.GRPCError(grpclib.const.Status.UNIMPLEMENTED)")
Comment on lines +126 to +141
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These are the lines I added compared to the grpclib version


buf.add("")
buf.add("")
buf.add("class {}Stub:", service.name)
with buf.indent():
buf.add("")
buf.add("def __init__(self, channel: {}.{}) -> None:".format(client.__name__, client.Channel.__name__))
with buf.indent():
if len(service.methods) == 0:
buf.add("pass")
for method in service.methods:
name, cardinality, request_type, reply_type = method
full_name = "/{}/{}".format(service_name, name)
method_cls: type
if cardinality is const.Cardinality.UNARY_UNARY:
method_cls = client.UnaryUnaryMethod
elif cardinality is const.Cardinality.UNARY_STREAM:
method_cls = client.UnaryStreamMethod
elif cardinality is const.Cardinality.STREAM_UNARY:
method_cls = client.StreamUnaryMethod
elif cardinality is const.Cardinality.STREAM_STREAM:
method_cls = client.StreamStreamMethod
else:
raise TypeError(cardinality)
method_cls = cast(type, method_cls) # FIXME: redundant
buf.add("self.{} = {}.{}(".format(name, client.__name__, method_cls.__name__))
with buf.indent():
buf.add("channel,")
buf.add("{!r},".format(full_name))
buf.add("{},", request_type)
buf.add("{},", reply_type)
buf.add(")")
return buf.content()


def _get_proto(request: CodeGeneratorRequest, name: str) -> FileDescriptorProto:
return next(f for f in request.proto_file if f.name == name)


def _strip_proto(proto_file_path: str) -> str:
for suffix in [".protodevel", ".proto"]:
if proto_file_path.endswith(suffix):
return proto_file_path[: -len(suffix)]

return proto_file_path


def _base_module_name(proto_file_path: str) -> str:
basename = _strip_proto(proto_file_path)
return basename.replace("-", "_").replace("/", ".")


def _proto2pb2_module_name(proto_file_path: str) -> str:
return _base_module_name(proto_file_path) + "_pb2"


def _proto2grpc_module_name(proto_file_path: str) -> str:
return _base_module_name(proto_file_path) + "_grpc"


def _type_names(
proto_file: FileDescriptorProto,
message_type: DescriptorProto,
parents: Optional[Deque[str]] = None,
) -> Iterator[Tuple[str, str]]:
if parents is None:
parents = deque()

proto_name_parts = [""]
if proto_file.package:
proto_name_parts.append(proto_file.package)
proto_name_parts.extend(parents)
proto_name_parts.append(message_type.name)

py_name_parts = [_proto2pb2_module_name(proto_file.name)]
py_name_parts.extend(parents)
py_name_parts.append(message_type.name)

yield ".".join(proto_name_parts), ".".join(py_name_parts)

parents.append(message_type.name)
for nested in message_type.nested_type:
yield from _type_names(proto_file, nested, parents=parents)
parents.pop()


def main() -> None:
with os.fdopen(sys.stdin.fileno(), "rb") as inp:
request = CodeGeneratorRequest.FromString(inp.read())

types_map: Dict[str, str] = {}
for pf in request.proto_file:
for mt in pf.message_type:
types_map.update(_type_names(pf, mt))

response = CodeGeneratorResponse()

# See https://github.com/protocolbuffers/protobuf/blob/v3.12.0/docs/implementing_proto3_presence.md # noqa
if hasattr(CodeGeneratorResponse, "Feature"):
response.supported_features = CodeGeneratorResponse.FEATURE_PROTO3_OPTIONAL

for file_to_generate in request.file_to_generate:
proto_file = _get_proto(request, file_to_generate)

imports = [_proto2pb2_module_name(dep) for dep in list(proto_file.dependency) + [file_to_generate]]

services = []
for service in proto_file.service:
methods = []
for method in service.method:
cardinality = _CARDINALITY[(method.client_streaming, method.server_streaming)] # type: ignore
methods.append(
Method(
name=method.name,
cardinality=cardinality,
request_type=types_map[method.input_type],
reply_type=types_map[method.output_type],
)
)
services.append(Service(name=service.name, methods=methods))

file = response.file.add()
module_name = _proto2grpc_module_name(file_to_generate)
file.name = module_name.replace(".", "/") + ".py"
file.content = render(
proto_file=proto_file.name,
package=proto_file.package,
imports=imports,
services=services,
)

with os.fdopen(sys.stdout.fileno(), "wb") as out:
out.write(response.SerializeToString())


if __name__ == "__main__":
sys.argv[0] = re.sub(r"(-script\.pyw|\.exe)?$", "", sys.argv[0])
sys.exit(main())
7 changes: 4 additions & 3 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,10 @@ include = ["LICENSE", "src/viam/rpc/libviam_rust_utils.*"]
]
pyright = "^1.1.339"

[tool.poetry.extras]
mlmodel = ["numpy"]


[tool.pytest.ini_options]
addopts = "-ra"
testpaths = "tests"
Expand All @@ -75,9 +79,6 @@ line_length = 140
requires = [ "poetry-core>=1.0.0" ]
build-backend = "poetry.core.masonry.api"

[tool.poetry.extras]
mlmodel = ["numpy"]

[tool.pyright]
include = [ "src" ]
exclude = [ "**/gen", "**/proto" ]
14 changes: 14 additions & 0 deletions src/viam/gen/app/agent/v1/agent_grpc.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
import typing
import grpclib.const
import grpclib.client
import grpclib.exceptions
if typing.TYPE_CHECKING:
import grpclib.server
import google.protobuf.duration_pb2
Expand All @@ -22,6 +23,14 @@ async def UpdateAgentConfig(self, stream: 'grpclib.server.Stream[app.agent.v1.ag
def __mapping__(self) -> typing.Dict[str, grpclib.const.Handler]:
return {'/viam.app.agent.v1.AgentAppService/GetAgentConfig': grpclib.const.Handler(self.GetAgentConfig, grpclib.const.Cardinality.UNARY_UNARY, app.agent.v1.agent_pb2.GetAgentConfigRequest, app.agent.v1.agent_pb2.GetAgentConfigResponse), '/viam.app.agent.v1.AgentAppService/UpdateAgentConfig': grpclib.const.Handler(self.UpdateAgentConfig, grpclib.const.Cardinality.UNARY_UNARY, app.agent.v1.agent_pb2.UpdateAgentConfigRequest, app.agent.v1.agent_pb2.UpdateAgentConfigResponse)}

class UnimplementedAgentAppServiceBase(AgentAppServiceBase):

async def GetAgentConfig(self, stream: 'grpclib.server.Stream[app.agent.v1.agent_pb2.GetAgentConfigRequest, app.agent.v1.agent_pb2.GetAgentConfigResponse]') -> None:
raise grpclib.exceptions.GRPCError(grpclib.const.Status.UNIMPLEMENTED)

async def UpdateAgentConfig(self, stream: 'grpclib.server.Stream[app.agent.v1.agent_pb2.UpdateAgentConfigRequest, app.agent.v1.agent_pb2.UpdateAgentConfigResponse]') -> None:
raise grpclib.exceptions.GRPCError(grpclib.const.Status.UNIMPLEMENTED)

class AgentAppServiceStub:

def __init__(self, channel: grpclib.client.Channel) -> None:
Expand All @@ -37,6 +46,11 @@ async def DeviceAgentConfig(self, stream: 'grpclib.server.Stream[app.agent.v1.ag
def __mapping__(self) -> typing.Dict[str, grpclib.const.Handler]:
return {'/viam.app.agent.v1.AgentDeviceService/DeviceAgentConfig': grpclib.const.Handler(self.DeviceAgentConfig, grpclib.const.Cardinality.UNARY_UNARY, app.agent.v1.agent_pb2.DeviceAgentConfigRequest, app.agent.v1.agent_pb2.DeviceAgentConfigResponse)}

class UnimplementedAgentDeviceServiceBase(AgentDeviceServiceBase):

async def DeviceAgentConfig(self, stream: 'grpclib.server.Stream[app.agent.v1.agent_pb2.DeviceAgentConfigRequest, app.agent.v1.agent_pb2.DeviceAgentConfigResponse]') -> None:
raise grpclib.exceptions.GRPCError(grpclib.const.Status.UNIMPLEMENTED)

class AgentDeviceServiceStub:

def __init__(self, channel: grpclib.client.Channel) -> None:
Expand Down
21 changes: 21 additions & 0 deletions src/viam/gen/app/cloudslam/v1/cloud_slam_grpc.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
import typing
import grpclib.const
import grpclib.client
import grpclib.exceptions
if typing.TYPE_CHECKING:
import grpclib.server
from .... import common
Expand Down Expand Up @@ -38,6 +39,26 @@ async def GetMappingSessionMetadataByID(self, stream: 'grpclib.server.Stream[app
def __mapping__(self) -> typing.Dict[str, grpclib.const.Handler]:
return {'/viam.app.cloudslam.v1.CloudSLAMService/StartMappingSession': grpclib.const.Handler(self.StartMappingSession, grpclib.const.Cardinality.UNARY_UNARY, app.cloudslam.v1.cloud_slam_pb2.StartMappingSessionRequest, app.cloudslam.v1.cloud_slam_pb2.StartMappingSessionResponse), '/viam.app.cloudslam.v1.CloudSLAMService/GetActiveMappingSessionsForRobot': grpclib.const.Handler(self.GetActiveMappingSessionsForRobot, grpclib.const.Cardinality.UNARY_UNARY, app.cloudslam.v1.cloud_slam_pb2.GetActiveMappingSessionsForRobotRequest, app.cloudslam.v1.cloud_slam_pb2.GetActiveMappingSessionsForRobotResponse), '/viam.app.cloudslam.v1.CloudSLAMService/GetMappingSessionPointCloud': grpclib.const.Handler(self.GetMappingSessionPointCloud, grpclib.const.Cardinality.UNARY_UNARY, app.cloudslam.v1.cloud_slam_pb2.GetMappingSessionPointCloudRequest, app.cloudslam.v1.cloud_slam_pb2.GetMappingSessionPointCloudResponse), '/viam.app.cloudslam.v1.CloudSLAMService/ListMappingSessions': grpclib.const.Handler(self.ListMappingSessions, grpclib.const.Cardinality.UNARY_UNARY, app.cloudslam.v1.cloud_slam_pb2.ListMappingSessionsRequest, app.cloudslam.v1.cloud_slam_pb2.ListMappingSessionsResponse), '/viam.app.cloudslam.v1.CloudSLAMService/StopMappingSession': grpclib.const.Handler(self.StopMappingSession, grpclib.const.Cardinality.UNARY_UNARY, app.cloudslam.v1.cloud_slam_pb2.StopMappingSessionRequest, app.cloudslam.v1.cloud_slam_pb2.StopMappingSessionResponse), '/viam.app.cloudslam.v1.CloudSLAMService/GetMappingSessionMetadataByID': grpclib.const.Handler(self.GetMappingSessionMetadataByID, grpclib.const.Cardinality.UNARY_UNARY, app.cloudslam.v1.cloud_slam_pb2.GetMappingSessionMetadataByIDRequest, app.cloudslam.v1.cloud_slam_pb2.GetMappingSessionMetadataByIDResponse)}

class UnimplementedCloudSLAMServiceBase(CloudSLAMServiceBase):

async def StartMappingSession(self, stream: 'grpclib.server.Stream[app.cloudslam.v1.cloud_slam_pb2.StartMappingSessionRequest, app.cloudslam.v1.cloud_slam_pb2.StartMappingSessionResponse]') -> None:
raise grpclib.exceptions.GRPCError(grpclib.const.Status.UNIMPLEMENTED)

async def GetActiveMappingSessionsForRobot(self, stream: 'grpclib.server.Stream[app.cloudslam.v1.cloud_slam_pb2.GetActiveMappingSessionsForRobotRequest, app.cloudslam.v1.cloud_slam_pb2.GetActiveMappingSessionsForRobotResponse]') -> None:
raise grpclib.exceptions.GRPCError(grpclib.const.Status.UNIMPLEMENTED)

async def GetMappingSessionPointCloud(self, stream: 'grpclib.server.Stream[app.cloudslam.v1.cloud_slam_pb2.GetMappingSessionPointCloudRequest, app.cloudslam.v1.cloud_slam_pb2.GetMappingSessionPointCloudResponse]') -> None:
raise grpclib.exceptions.GRPCError(grpclib.const.Status.UNIMPLEMENTED)

async def ListMappingSessions(self, stream: 'grpclib.server.Stream[app.cloudslam.v1.cloud_slam_pb2.ListMappingSessionsRequest, app.cloudslam.v1.cloud_slam_pb2.ListMappingSessionsResponse]') -> None:
raise grpclib.exceptions.GRPCError(grpclib.const.Status.UNIMPLEMENTED)

async def StopMappingSession(self, stream: 'grpclib.server.Stream[app.cloudslam.v1.cloud_slam_pb2.StopMappingSessionRequest, app.cloudslam.v1.cloud_slam_pb2.StopMappingSessionResponse]') -> None:
raise grpclib.exceptions.GRPCError(grpclib.const.Status.UNIMPLEMENTED)

async def GetMappingSessionMetadataByID(self, stream: 'grpclib.server.Stream[app.cloudslam.v1.cloud_slam_pb2.GetMappingSessionMetadataByIDRequest, app.cloudslam.v1.cloud_slam_pb2.GetMappingSessionMetadataByIDResponse]') -> None:
raise grpclib.exceptions.GRPCError(grpclib.const.Status.UNIMPLEMENTED)

class CloudSLAMServiceStub:

def __init__(self, channel: grpclib.client.Channel) -> None:
Expand Down
Loading
Loading