Ver Fonte

Add protoc-gen-chattolib plugin that emits async+sync Connect client stubs

the-codemancer há 2 meses atrás
pai
commit
e6cdc4b10e
3 ficheiros alterados com 177 adições e 1 exclusões
  1. 3 1
      pyproject.toml
  2. 122 0
      src/chattolib/_protoc_plugin.py
  3. 52 0
      tests/test_protoc_plugin.py

+ 3 - 1
pyproject.toml

@@ -25,13 +25,15 @@ classifiers = [
     "Framework :: AsyncIO",
 ]
 dependencies = [
-    "connectrpc>=0.11",
     "protobuf>=5.28",
 ]
 
 [project.urls]
 Homepage = "https://chat.chatto.run"
 
+[project.scripts]
+protoc-gen-chattolib = "chattolib._protoc_plugin:main"
+
 [project.optional-dependencies]
 realtime = [
     "websockets>=13.0",

+ 122 - 0
src/chattolib/_protoc_plugin.py

@@ -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()

+ 52 - 0
tests/test_protoc_plugin.py

@@ -0,0 +1,52 @@
+"""The plugin turns a CodeGeneratorRequest into async+sync client stubs."""
+
+from chattolib._pb.google.protobuf import protoc_gen_request_pb2 as g
+from chattolib._protoc_plugin import generate
+
+
+def _request_with_room_service() -> g.CodeGeneratorRequest:
+    req = g.CodeGeneratorRequest()
+    req.file_to_generate.append("chatto/api/v1/rooms.proto")
+    fd = req.proto_file.add()
+    fd.name = "chatto/api/v1/rooms.proto"
+    fd.package = "chatto.api.v1"
+    svc = fd.service.add()
+    svc.name = "RoomService"
+    for name in ("JoinRoom", "CreateRoom"):
+        m = svc.method.add()
+        m.name = name
+        m.input_type = f".chatto.api.v1.{name}Request"
+        m.output_type = f".chatto.api.v1.{name}Response"
+    return req
+
+
+def test_emits_one_file_per_service_with_both_clients():
+    resp = generate(_request_with_room_service())
+    assert len(resp.file) == 1
+    out = resp.file[0]
+    assert out.name.endswith("rooms_connect.py")
+    content = out.content
+    # Both client classes, correct base classes, one method per RPC.
+    assert "class RoomServiceClient(ConnectClient):" in content
+    assert "class RoomServiceClientSync(ConnectClientSync):" in content
+    assert "async def join_room(" in content
+    assert "async def create_room(" in content
+    assert "def join_room(" in content  # sync variant
+    # Imports only from chattolib._connect and the sibling pb2.
+    assert "from chattolib._connect import ConnectClient, ConnectClientSync, MethodInfo" in content
+    assert "import chatto.api.v1.rooms_pb2 as _pb2" in content
+    # No connectrpc anywhere.
+    assert "connectrpc" not in content
+    # use_get is dropped.
+    assert "use_get" not in content
+
+
+def test_skips_files_without_services():
+    req = g.CodeGeneratorRequest()
+    req.file_to_generate.append("chatto/realtime/v1/realtime.proto")
+    fd = req.proto_file.add()
+    fd.name = "chatto/realtime/v1/realtime.proto"
+    fd.package = "chatto.realtime.v1"
+    fd.message_type.add().name = "RealtimeClientFrame"  # messages, no service
+    resp = generate(req)
+    assert len(resp.file) == 0