|
@@ -0,0 +1,122 @@
|
|
|
|
|
+"""A `protoc` plugin that generates chattolib's Connect client stubs.
|
|
|
|
|
+
|
|
|
|
|
+Speaks the standard plugin protocol: reads a ``CodeGeneratorRequest`` from
|
|
|
|
|
+stdin, writes a ``CodeGeneratorResponse`` to stdout. For each proto that
|
|
|
|
|
+declares a ``service``, it emits one ``<name>_connect.py`` defining an async
|
|
|
|
|
+``*ServiceClient(ConnectClient)`` and a sync ``*ServiceClientSync
|
|
|
|
|
+(ConnectClientSync)`` — both delegating to the shared runtime in
|
|
|
|
|
+``chattolib._connect``.
|
|
|
|
|
+
|
|
|
|
|
+Installed as the console script ``protoc-gen-chattolib`` so ``protoc`` finds
|
|
|
|
|
+it on PATH and it is invoked with ``--chattolib_out=<dir>``.
|
|
|
|
|
+"""
|
|
|
|
|
+
|
|
|
|
|
+from __future__ import annotations
|
|
|
|
|
+
|
|
|
|
|
+import sys
|
|
|
|
|
+from typing import Any
|
|
|
|
|
+
|
|
|
|
|
+from chattolib._pb.google.protobuf import (
|
|
|
|
|
+ protoc_gen_request_pb2 as g,
|
|
|
|
|
+)
|
|
|
|
|
+from chattolib._pb.google.protobuf import (
|
|
|
|
|
+ protoc_gen_response_pb2 as r,
|
|
|
|
|
+)
|
|
|
|
|
+
|
|
|
|
|
+_HEADER = "# Generated by chattolib's protoc plugin (protoc-gen-chattolib). DO NOT EDIT!\n"
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+def _method_name(m: Any) -> str:
|
|
|
|
|
+ """snake_case the RPC name, matching Python method naming."""
|
|
|
|
|
+ out = []
|
|
|
|
|
+ for i, ch in enumerate(m.name):
|
|
|
|
|
+ if ch.isupper() and i:
|
|
|
|
|
+ out.append("_")
|
|
|
|
|
+ out.append(ch.lower())
|
|
|
|
|
+ return "".join(out)
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+def _service_stub(fd: Any) -> str:
|
|
|
|
|
+ """Render one service's async + sync client classes as source text."""
|
|
|
|
|
+ pkg = fd.package
|
|
|
|
|
+ svc = fd.service[0]
|
|
|
|
|
+ base = f"{pkg}.{svc.name}" if pkg else svc.name
|
|
|
|
|
+ # The sibling message module is the proto's path with the .proto suffix
|
|
|
|
|
+ # turned into a dotted import path (e.g. chatto/api/v1/rooms.proto ->
|
|
|
|
|
+ # chatto.api.v1.rooms_pb2), resolved via the _pb sys.path shim.
|
|
|
|
|
+ pb2_module = fd.name[: -len(".proto")].replace("/", ".") + "_pb2"
|
|
|
|
|
+ lines = [
|
|
|
|
|
+ _HEADER,
|
|
|
|
|
+ f"# source: {fd.name}",
|
|
|
|
|
+ "",
|
|
|
|
|
+ f"import {pb2_module} as _pb2",
|
|
|
|
|
+ "from chattolib._connect import ConnectClient, ConnectClientSync, MethodInfo",
|
|
|
|
|
+ "",
|
|
|
|
|
+ "",
|
|
|
|
|
+ f"class {svc.name}Client(ConnectClient):",
|
|
|
|
|
+ ]
|
|
|
|
|
+ for m in svc.method:
|
|
|
|
|
+ mn = _method_name(m)
|
|
|
|
|
+ in_name = m.input_type.lstrip(".").split(".")[-1]
|
|
|
|
|
+ out_name = m.output_type.lstrip(".").split(".")[-1]
|
|
|
|
|
+ lines += [
|
|
|
|
|
+ f" async def {mn}(self, request: _pb2.{in_name}, *,"
|
|
|
|
|
+ f" headers: dict[str, str] | None = None) -> _pb2.{out_name}:",
|
|
|
|
|
+ " return await self.execute_unary(",
|
|
|
|
|
+ " request=request,",
|
|
|
|
|
+ " method=MethodInfo(",
|
|
|
|
|
+ f' name="{m.name}",',
|
|
|
|
|
+ f' service_name="{base}",',
|
|
|
|
|
+ f" input=_pb2.{in_name},",
|
|
|
|
|
+ f" output=_pb2.{out_name},",
|
|
|
|
|
+ " ),",
|
|
|
|
|
+ " headers=headers,",
|
|
|
|
|
+ " )",
|
|
|
|
|
+ "",
|
|
|
|
|
+ ]
|
|
|
|
|
+ lines += [f"class {svc.name}ClientSync(ConnectClientSync):", ""]
|
|
|
|
|
+ for m in svc.method:
|
|
|
|
|
+ mn = _method_name(m)
|
|
|
|
|
+ in_name = m.input_type.lstrip(".").split(".")[-1]
|
|
|
|
|
+ out_name = m.output_type.lstrip(".").split(".")[-1]
|
|
|
|
|
+ lines += [
|
|
|
|
|
+ f" def {mn}(self, request: _pb2.{in_name}, *,"
|
|
|
|
|
+ f" headers: dict[str, str] | None = None) -> _pb2.{out_name}:",
|
|
|
|
|
+ " return self.execute_unary(",
|
|
|
|
|
+ " request=request,",
|
|
|
|
|
+ " method=MethodInfo(",
|
|
|
|
|
+ f' name="{m.name}",',
|
|
|
|
|
+ f' service_name="{base}",',
|
|
|
|
|
+ f" input=_pb2.{in_name},",
|
|
|
|
|
+ f" output=_pb2.{out_name},",
|
|
|
|
|
+ " ),",
|
|
|
|
|
+ " headers=headers,",
|
|
|
|
|
+ " )",
|
|
|
|
|
+ "",
|
|
|
|
|
+ ]
|
|
|
|
|
+ return "\n".join(lines)
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+def generate(request: g.CodeGeneratorRequest) -> r.CodeGeneratorResponse:
|
|
|
|
|
+ resp = r.CodeGeneratorResponse()
|
|
|
|
|
+ # file_to_generate is a repeated *string* (proto file names), not a list of
|
|
|
|
|
+ # message objects — match by string membership.
|
|
|
|
|
+ to_generate = set(request.file_to_generate)
|
|
|
|
|
+ for fd in request.proto_file:
|
|
|
|
|
+ if not fd.service or fd.name not in to_generate:
|
|
|
|
|
+ continue
|
|
|
|
|
+ out = resp.file.add()
|
|
|
|
|
+ out.name = fd.name.replace(".proto", "_connect.py")
|
|
|
|
|
+ out.content = _service_stub(fd)
|
|
|
|
|
+ return resp
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+def main() -> None:
|
|
|
|
|
+ data = sys.stdin.buffer.read()
|
|
|
|
|
+ request = g.CodeGeneratorRequest.FromString(data)
|
|
|
|
|
+ resp = generate(request)
|
|
|
|
|
+ sys.stdout.buffer.write(resp.SerializeToString())
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+if __name__ == "__main__":
|
|
|
|
|
+ main()
|