Skip to content

Commit ce7302a

Browse files
zcheezchee
committed
WIP
Signed-off-by: Koichi Shiraishi <[email protected]>
1 parent a4765b0 commit ce7302a

File tree

3 files changed

+89
-39
lines changed

3 files changed

+89
-39
lines changed

‎cmd/protoc-gen-connect-python/generator/generator.go‎

Lines changed: 50 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -20,8 +20,6 @@ import (
2020
"google.golang.org/protobuf/types/pluginpb"
2121
)
2222

23-
var version = "devel"
24-
2523
type Config struct{}
2624

2725
type Generator struct {
@@ -103,6 +101,7 @@ type Method struct {
103101
Method string
104102
FullName string
105103
RPCType RPCType
104+
Options protoreflect.ProtoMessage
106105
}
107106

108107
type message struct {
@@ -140,16 +139,17 @@ func (g *Generator) generate(gen *protogen.GeneratedFile, f *protogen.File) {
140139
idx: i,
141140
Method: meth.GoName,
142141
FullName: fullname[:idx],
142+
Options: meth.Desc.Options(),
143143
}
144144

145145
// parse RPC type
146146
switch {
147+
case meth.Desc.IsStreamingServer() && meth.Desc.IsStreamingClient():
148+
method.RPCType = BidirectionalStreaming
147149
case meth.Desc.IsStreamingServer():
148150
method.RPCType = ServerStreaming
149151
case meth.Desc.IsStreamingClient():
150152
method.RPCType = ClientStreaming
151-
case meth.Desc.IsStreamingServer() && meth.Desc.IsStreamingClient():
152-
method.RPCType = BidirectionalStreaming
153153
default:
154154
method.RPCType = Unary
155155
}
@@ -174,13 +174,23 @@ func (g *Generator) generate(gen *protogen.GeneratedFile, f *protogen.File) {
174174
}
175175
p.P(`"""Generated connect code."""`)
176176
p.P()
177+
p.P(`import abc`)
177178
p.P(`from enum import Enum`)
178179
p.P()
179-
p.P(`from connect.client import Client`)
180-
p.P(`from connect.connect import StreamRequest, StreamResponse, UnaryRequest, UnaryResponse`)
181-
p.P(`from connect.handler import ClientStreamHandler, Handler, ServerStreamHandler, UnaryHandler`)
182-
p.P(`from connect.options import ClientOptions, ConnectOptions`)
180+
p.P(`from connect import (`)
181+
p.P(` Client,`)
182+
p.P(` ClientOptions,`)
183+
p.P(` ConnectOptions,`)
184+
p.P(` Handler,`)
185+
p.P(` HandlerContext,`)
186+
p.P(` IdempotencyLevel,`)
187+
p.P(` StreamRequest,`)
188+
p.P(` StreamResponse,`)
189+
p.P(` UnaryRequest,`)
190+
p.P(` UnaryResponse,`)
191+
p.P(`)`)
183192
p.P(`from connect.connection_pool import AsyncConnectionPool`)
193+
p.P(`from connect.handler import BidiStreamHandler, ClientStreamHandler, ServerStreamHandler, UnaryHandler`)
184194
p.P(`from google.protobuf.descriptor import MethodDescriptor, ServiceDescriptor`)
185195
p.P()
186196

@@ -191,7 +201,7 @@ func (g *Generator) generate(gen *protogen.GeneratedFile, f *protogen.File) {
191201
var sb strings.Builder
192202
numSvc := len(p.services)
193203
if numSvc > 0 {
194-
fmt.Fprintf(&sb, "from ..%s import ", svcNamePB)
204+
fmt.Fprintf(&sb, "from ..%s import (\n", svcNamePB)
195205
}
196206
seem := make(map[string]bool)
197207
i := 0
@@ -201,22 +211,25 @@ func (g *Generator) generate(gen *protogen.GeneratedFile, f *protogen.File) {
201211
case seem[svc.input.method] && seem[svc.output.method]:
202212
continue
203213
case !seem[svc.input.method] && seem[svc.output.method]:
204-
fmt.Fprintf(&sb, "%s", svc.input.method)
214+
fmt.Fprintf(&sb, " %s\n", svc.input.method)
205215
seem[svc.input.method] = true
206216
case seem[svc.input.method] && !seem[svc.output.method]:
207-
fmt.Fprintf(&sb, "%s", svc.output.method)
217+
fmt.Fprintf(&sb, " %s\n", svc.output.method)
208218
seem[svc.output.method] = true
209219
default:
210-
fmt.Fprintf(&sb, "%s, %s", svc.input.method, svc.output.method)
220+
fmt.Fprintf(&sb, " %s,\n %s", svc.input.method, svc.output.method)
211221
seem[svc.input.method] = true
212222
seem[svc.output.method] = true
213223
}
214224
if i <= numSvc-2 {
215-
fmt.Fprint(&sb, ", ")
225+
fmt.Fprint(&sb, ",\n")
226+
} else {
227+
fmt.Fprint(&sb, ",")
216228
}
217229
i++
218230
}
219-
p.P(strings.TrimSuffix(sb.String(), ", "))
231+
p.P(sb.String())
232+
p.P(`)`)
220233
p.P()
221234
p.P()
222235
procedures := upperSvcName + `Procedures`
@@ -244,20 +257,28 @@ func (g *Generator) generate(gen *protogen.GeneratedFile, f *protogen.File) {
244257
for _, meth := range sortedMap(p.services) {
245258
svc := p.services[meth]
246259
p.P(` `, `self.`, meth.Method, ` = `, `Client[`, svc.input.method, `, `, svc.output.method, `](`)
247-
p.P(` `, `pool, `, `base_url + `, procedures+`.`+meth.Method+`.value, `, svc.input.method+`, `, svc.output.method, `, options`)
260+
if options := meth.Options; options != nil {
261+
if desc, ok := options.(*descriptorpb.MethodOptions); ok && desc.GetIdempotencyLevel() != descriptorpb.MethodOptions_IDEMPOTENCY_UNKNOWN {
262+
p.P(` `, `pool, `, `base_url + `, procedures+`.`+meth.Method+`.value, `, svc.input.method+`, `, svc.output.method, `, ClientOptions(idempotency_level=IdempotencyLevel.NO_SIDE_EFFECTS, enable_get=True).merge(options)`)
263+
} else {
264+
p.P(` `, `pool, `, `base_url + `, procedures+`.`+meth.Method+`.value, `, svc.input.method+`, `, svc.output.method, `, options`)
265+
}
266+
}
248267
switch meth.RPCType {
249268
case Unary:
250269
p.P(` `, `).call_unary`)
251270
case ServerStreaming:
252271
p.P(` `, `).call_server_stream`)
253272
case ClientStreaming:
254273
p.P(` `, `).call_client_stream`)
274+
case BidirectionalStreaming:
275+
p.P(` `, `).call_bidi_stream`)
255276
}
256277
}
257278
p.P()
258279
p.P()
259280
handler := upperSvcName + `Handler`
260-
p.P(`class `, handler, `:`)
281+
p.P(`class `, handler, `(metaclass=abc.ABCMeta):`)
261282
p.P(` `, `"""Handler for the `, lowerCamelCase(upperSvcName), ` service."""`)
262283
p.P()
263284
j := 0
@@ -269,16 +290,15 @@ func (g *Generator) generate(gen *protogen.GeneratedFile, f *protogen.File) {
269290
reqRPCType string
270291
respRPCType string
271292
)
272-
// TODO(zchee): BidirectionalStreaming?
273293
switch meth.RPCType {
274294
case Unary:
275295
reqRPCType = `UnaryRequest`
276296
respRPCType = `UnaryResponse`
277-
case ServerStreaming, ClientStreaming:
297+
case ServerStreaming, ClientStreaming, BidirectionalStreaming:
278298
reqRPCType = `StreamRequest`
279299
respRPCType = `StreamResponse`
280300
}
281-
fmt.Fprintf(&sb, "%s[%s]) -> %s[%s]: ...", reqRPCType, svc.input.method, respRPCType, svc.output.method)
301+
fmt.Fprintf(&sb, "%s[%s], context: HandlerContext) -> %s[%s]:\n raise NotImplementedError()", reqRPCType, svc.input.method, respRPCType, svc.output.method)
282302
p.P(sb.String())
283303
if j <= len(p.services)-2 {
284304
p.P()
@@ -295,7 +315,6 @@ func (g *Generator) generate(gen *protogen.GeneratedFile, f *protogen.File) {
295315
rpcHandler string
296316
call string
297317
)
298-
// TODO(zchee): BidirectionalStreaming?
299318
switch meth.RPCType {
300319
case Unary:
301320
rpcHandler = `UnaryHandler`
@@ -306,17 +325,25 @@ func (g *Generator) generate(gen *protogen.GeneratedFile, f *protogen.File) {
306325
case ClientStreaming:
307326
rpcHandler = `ClientStreamHandler`
308327
call = fmt.Sprintf(" stream=service.%s,", meth.Method)
328+
case BidirectionalStreaming:
329+
rpcHandler = `BidiStreamHandler`
330+
call = fmt.Sprintf(" stream=service.%s,", meth.Method)
309331
}
310332

311-
// TODO(zchee): BidirectionalStreaming?
312333
switch meth.RPCType {
313-
case Unary, ServerStreaming, ClientStreaming:
334+
case Unary, ServerStreaming, ClientStreaming, BidirectionalStreaming:
314335
p.P(` `, rpcHandler, `(`)
315336
p.P(` procedure=`, procedures+`.`+meth.Method+`.value,`)
316337
p.P(call)
317338
p.P(` input=`, svc.input.method, `,`)
318339
p.P(` output=`, svc.output.method, `,`)
319-
p.P(` options=options,`)
340+
if options := meth.Options; options != nil {
341+
if desc, ok := options.(*descriptorpb.MethodOptions); ok && desc.GetIdempotencyLevel() != descriptorpb.MethodOptions_IDEMPOTENCY_UNKNOWN {
342+
p.P(` options=ConnectOptions(idempotency_level=IdempotencyLevel.NO_SIDE_EFFECTS).merge(options),`)
343+
} else {
344+
p.P(` options=options,`)
345+
}
346+
}
320347
p.P(` ),`)
321348
}
322349
}

‎examples/proto/connectrpc/eliza/v1/v1connect/eliza_connect.py‎

Lines changed: 38 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,20 +1,39 @@
11
# Generated by the protoc-gen-connect-python. DO NOT EDIT!
22
# source: examples/proto/connectrpc/eliza/v1/v1connect/eliza.proto
33
# Protobuf Python Version: v5.29.3
4-
# protoc-gen-connect-python version: v0.0.0-20250517015031-b19a36b52499+dirty
4+
# protoc-gen-connect-python version: v0.0.0-20250707225833-133e5e81d930
55
"""Generated connect code."""
66

7+
import abc
78
from enum import Enum
89

9-
from connect.client import Client
10-
from connect.connect import StreamRequest, StreamResponse, UnaryRequest, UnaryResponse
11-
from connect.handler import ClientStreamHandler, Handler, ServerStreamHandler, UnaryHandler
12-
from connect.options import ClientOptions, ConnectOptions
10+
from connect import (
11+
Client,
12+
ClientOptions,
13+
ConnectOptions,
14+
Handler,
15+
HandlerContext,
16+
IdempotencyLevel,
17+
StreamRequest,
18+
StreamResponse,
19+
UnaryRequest,
20+
UnaryResponse,
21+
)
1322
from connect.connection_pool import AsyncConnectionPool
23+
from connect.handler import BidiStreamHandler, ClientStreamHandler, ServerStreamHandler, UnaryHandler
1424
from google.protobuf.descriptor import MethodDescriptor, ServiceDescriptor
1525

1626
from .. import eliza_pb2
17-
from ..eliza_pb2 import SayRequest, SayResponse, ConverseRequest, ConverseResponse, IntroduceRequest, IntroduceResponse, ReflectRequest, ReflectResponse
27+
from ..eliza_pb2 import (
28+
SayRequest,
29+
SayResponse,
30+
ConverseRequest,
31+
ConverseResponse,
32+
IntroduceRequest,
33+
IntroduceResponse,
34+
ReflectRequest,
35+
ReflectResponse,
36+
)
1837

1938

2039
class ElizaServiceProcedures(Enum):
@@ -39,11 +58,11 @@ def __init__(self, base_url: str, pool: AsyncConnectionPool, options: ClientOpti
3958
base_url = base_url.removesuffix("/")
4059

4160
self.Say = Client[SayRequest, SayResponse](
42-
pool, base_url + ElizaServiceProcedures.Say.value, SayRequest, SayResponse, options
61+
pool, base_url + ElizaServiceProcedures.Say.value, SayRequest, SayResponse, ClientOptions(idempotency_level=IdempotencyLevel.NO_SIDE_EFFECTS, enable_get=True).merge(options)
4362
).call_unary
4463
self.Converse = Client[ConverseRequest, ConverseResponse](
4564
pool, base_url + ElizaServiceProcedures.Converse.value, ConverseRequest, ConverseResponse, options
46-
).call_server_stream
65+
).call_bidi_stream
4766
self.Introduce = Client[IntroduceRequest, IntroduceResponse](
4867
pool, base_url + ElizaServiceProcedures.Introduce.value, IntroduceRequest, IntroduceResponse, options
4968
).call_server_stream
@@ -52,16 +71,20 @@ def __init__(self, base_url: str, pool: AsyncConnectionPool, options: ClientOpti
5271
).call_client_stream
5372

5473

55-
class ElizaServiceHandler:
74+
class ElizaServiceHandler(metaclass=abc.ABCMeta):
5675
"""Handler for the elizaService service."""
5776

58-
async def Say(self, request: UnaryRequest[SayRequest]) -> UnaryResponse[SayResponse]: ...
77+
async def Say(self, request: UnaryRequest[SayRequest], context: HandlerContext) -> UnaryResponse[SayResponse]:
78+
raise NotImplementedError()
5979

60-
async def Converse(self, request: StreamRequest[ConverseRequest]) -> StreamResponse[ConverseResponse]: ...
80+
async def Converse(self, request: StreamRequest[ConverseRequest], context: HandlerContext) -> StreamResponse[ConverseResponse]:
81+
raise NotImplementedError()
6182

62-
async def Introduce(self, request: StreamRequest[IntroduceRequest]) -> StreamResponse[IntroduceResponse]: ...
83+
async def Introduce(self, request: StreamRequest[IntroduceRequest], context: HandlerContext) -> StreamResponse[IntroduceResponse]:
84+
raise NotImplementedError()
6385

64-
async def Reflect(self, request: StreamRequest[ReflectRequest]) -> StreamResponse[ReflectResponse]: ...
86+
async def Reflect(self, request: StreamRequest[ReflectRequest], context: HandlerContext) -> StreamResponse[ReflectResponse]:
87+
raise NotImplementedError()
6588

6689

6790
def create_ElizaService_handlers(service: ElizaServiceHandler, options: ConnectOptions | None = None) -> list[Handler]:
@@ -71,9 +94,9 @@ def create_ElizaService_handlers(service: ElizaServiceHandler, options: ConnectO
7194
unary=service.Say,
7295
input=SayRequest,
7396
output=SayResponse,
74-
options=options,
97+
options=ConnectOptions(idempotency_level=IdempotencyLevel.NO_SIDE_EFFECTS).merge(options),
7598
),
76-
ServerStreamHandler(
99+
BidiStreamHandler(
77100
procedure=ElizaServiceProcedures.Converse.value,
78101
stream=service.Converse,
79102
input=ConverseRequest,

‎examples/proto/connectrpc/eliza/v1/v1connect/eliza_connect_pb2.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -58,7 +58,7 @@ def __init__(self, base_url: str, pool: AsyncConnectionPool, options: ClientOpti
5858
base_url = base_url.removesuffix("/")
5959

6060
self.Say = Client[SayRequest, SayResponse](
61-
pool, base_url + ElizaServiceProcedures.Say.value, SayRequest, SayResponse, options
61+
pool, base_url + ElizaServiceProcedures.Say.value, SayRequest, SayResponse, ClientOptions(idempotency_level=IdempotencyLevel.NO_SIDE_EFFECTS, enable_get=True).merge(options)
6262
).call_unary
6363
self.Converse = Client[ConverseRequest, ConverseResponse](
6464
pool, base_url + ElizaServiceProcedures.Converse.value, ConverseRequest, ConverseResponse, options

0 commit comments

Comments
 (0)