Skip to content

Commit eade29c

Browse files
committed
Add first-class Python gRPC support (servicer binding + client detection)
enola already produced gRPC server routes from .proto files (language-agnostic) and detected the client side + bound server routes to handlers for Go and TypeScript. Python had neither: a servicer was a plain class and a stub.Method() call produced no route, so Python dropped out of cross-repo gRPC graphs, impact analysis, and unused-route detection. This brings Python to parity. Detection works from hand-written source + the .proto, not from generated *_pb2_grpc.py stubs — real repos (vosk-server, airflow) build those at build time and never commit them, so scanning generated stubs finds nothing. Client side: - New pythonextractor/grpcpy.go detects `stub = mod._pb2_grpc.FooStub(channel)` bindings and `stub.Method()` calls, emitting client-role routes. Binding is positional (interleaved by source offset), so a stub var rebound to a second service — the vosk pattern — resolves each call to the right service. - Python source only knows the short service name, so routes are emitted with a provisional short Name; a new engine pass resolvePyGRPCClientRoutes rewrites them to the fully-qualified wire path (from the proto's server routes) BEFORE linkCrossRepo, which matches routes by Name. Remove-then-add keeps the store's name index consistent. Unresolved (no proto / ambiguous short name) routes are left provisional. Server side: - Generalize bindGRPCHandlers/implShortName to bind a proto server route to a Python servicer method via the `class X(mod._pb2_grpc.FooServicer)` convention (in addition to Go's Unimplemented<Service>Server embed), indexing Python classes alongside Go structs. The existing per-repo scoping and ambiguity guard are unchanged — two impls of one service in a repo (vosk grpc/ + grpc-wav2vec/) correctly stay unbound. Also: classify python-grpc-client in crossrepo handWrittenClientSources; bump cacheVersion v87 -> v88 (Python extractor now emits new facts); register cachecov coverage for v88 and for v87 (the resolveCall fix, whose registration was missing). Tests: unit tests for client detection (incl. positional rebinding + adversarial GrpcHook/no-import negatives) and for servicer binding + FQ resolution; a golden py_grpc_multirepo fixture (trimmed from real vosk-server) exercising the full pipeline — proto server routes, handler binding, client FQ resolution, cross-repo link, and an unmatched RPC. Verified live via MCP against real vosk-server and against airflow (no regression).
1 parent c8fb939 commit eade29c

16 files changed

Lines changed: 589 additions & 7 deletions

File tree

internal/cachecov/coverage_test.go

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -116,8 +116,10 @@ var versionCoverage = map[int][]string{
116116
82: {"TestAST_RegistrationDecoratorsMarkUsed", "TestAST_AttributeArgValueRef", "TestAST_DictValueRef", "TestAST_ListOfLocalsNoFalseRef"}, // Python registration decorators + attribute/collection value-refs
117117
83: {"TestPyComplexity_ScalingLoopDepth_BoundedInnerDiscounted", "TestPyComplexity_ScalingLoopDepth_FullyBounded", "TestPyComplexity_ScalingLoopDepth_WhileTrueBounded", "TestPyComplexity_IODirect", "TestPyComputePerformsIO_Transitive", "TestExtract_ScalingLoopDepth_BoundedDiscounted"}, // scaling_loop_depth (Python/Go) + Python io_direct/performs_io
118118
84: {"TestPyComplexity_CallsInScalingLoop_BoundedExcluded", "TestExtract_CallsInScalingLoop_BoundedExcluded", "TestTsComplexity_CallsInScalingLoop_BoundedExcluded"}, // calls_in_scaling_loop (Python/Go/TS)
119-
85: {"TestAST_DataClassAndEnumProps", "TestAST_InformalAbstractDetection"}, // Python enum/data_class props + informal-abstract (raise NotImplementedError)
120-
86: {"TestAST_DataClassAndEnumProps"}, // Python data_class broadened to RootModel/*BaseModel subclasses (StrictBaseModel)
119+
85: {"TestAST_DataClassAndEnumProps", "TestAST_InformalAbstractDetection"}, // Python enum/data_class props + informal-abstract (raise NotImplementedError)
120+
86: {"TestAST_DataClassAndEnumProps"}, // Python data_class broadened to RootModel/*BaseModel subclasses (StrictBaseModel)
121+
87: {"TestAST_ParamCall_NoEdge", "TestAST_LocalCallable_NoEdge", "TestAST_LoopVarCall_NoEdge", "TestAST_SameModuleCall_StillResolves"}, // Python resolveCall no longer fabricates same-module edges for params/locals/loop vars
122+
88: {"TestPyGRPC_ClientStubCall_EmitsRoute", "TestPyGRPC_StubRebinding_PositionalBinding", "TestPyGRPC_DynamicStubClass_NoRoute", "TestPyGRPC_NoStubImport_NoRoute"}, // Python gRPC client-role routes from stub.Method() call sites
121123
}
122124

123125
func TestCacheVersionCoverage(t *testing.T) {

internal/engine/cache.go

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -226,7 +226,10 @@ import (
226226
// parameters, locals, and loop variables — it now resolves a bare callee only when the name
227227
// is a known module-level def and not shadowed by a parameter (mirrors valueRefTarget). Drops
228228
// spurious call edges; bump so cached Python snapshots re-extract with the tightened resolver.
229-
const cacheVersion = "v87"
229+
// v88: Python extractor now emits gRPC client-role routes (source=python-grpc-client) for
230+
// stub.Method(...) call sites, detected from source. New facts, so cached Python snapshots must
231+
// re-extract to pick them up.
232+
const cacheVersion = "v88"
230233

231234
// extractorCache holds per-extractor facts keyed by a content hash of the files
232235
// the extractor depends on. It is loaded from disk at the start of a snapshot and

internal/engine/engine.go

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -245,6 +245,7 @@ func (e *Engine) GenerateSnapshot(ctx context.Context, repoPath string, appendMo
245245
// import/shared-lib references. Recomputed from scratch each run (prior
246246
// synthetic facts are dropped first) so it stays idempotent across appends.
247247
tStage = time.Now()
248+
e.resolvePyGRPCClientRoutes()
248249
e.linkCrossRepo()
249250
e.flagUnmatchedRoutes()
250251
e.bindGRPCHandlers()

internal/engine/golden_test.go

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,7 @@ var fixtures = []fixture{
5858
{name: "multirepo", subRepos: []string{"repoA", "repoB"}},
5959
{name: "php_multirepo", subRepos: []string{"provider", "consumer"}},
6060
{name: "go_grpc_multirepo", subRepos: []string{"server", "client"}},
61+
{name: "py_grpc_multirepo", subRepos: []string{"server", "client"}},
6162
}
6263

6364
func TestGolden(t *testing.T) {

internal/engine/grpcbind.go

Lines changed: 100 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,12 @@ import (
1414
// alias-qualified) — capturing the service short name ("UserService").
1515
var unimplementedEmbed = regexp.MustCompile(`^(?:.*\.)?Unimplemented(.+)Server$`)
1616

17+
// servicerBase matches the target of the `implements` edge a Python gRPC servicer
18+
// carries by subclassing its generated base — e.g.
19+
// "stt_service_pb2_grpc.SttServiceServicer" — capturing the service short name
20+
// ("SttService"). grpc_python_out always names the base "<Service>Servicer".
21+
var servicerBase = regexp.MustCompile(`^(?:.*\.)?(.+)Servicer$`)
22+
1723
// bindGRPCHandlers connects each gRPC server route (emitted from a .proto by the
1824
// grpc extractor) to the Go method that implements it, so route → handler is
1925
// traversable by impact_analysis and find_path.
@@ -41,7 +47,7 @@ func (e *Engine) bindGRPCHandlers() {
4147
for _, s := range symbols {
4248
kind, _ := s.Props["symbol_kind"].(string)
4349
switch kind {
44-
case facts.SymbolStruct:
50+
case facts.SymbolStruct, facts.SymbolClass:
4551
short := implShortName(s)
4652
if short == "" {
4753
continue
@@ -99,8 +105,96 @@ func (e *Engine) bindGRPCHandlers() {
99105
}
100106
}
101107

102-
// implShortName returns the gRPC service short name a struct implements by
103-
// embedding Unimplemented<Service>Server, or "" if it embeds no such type.
108+
// resolvePyGRPCClientRoutes rewrites provisional Python gRPC client routes to their
109+
// fully-qualified wire path. The Python extractor only sees the short service name
110+
// (e.g. "SttService"), so it emits Name = "/SttService/StreamingRecognize" plus a
111+
// grpc_service_short prop; the fully-qualified name ("vosk.stt.v1.SttService") lives
112+
// only in the .proto. This pass resolves short → fq from the gRPC server routes in
113+
// the store and rewrites the Name to "/vosk.stt.v1.SttService/StreamingRecognize"
114+
// so it matches the server route at cross-repo link time.
115+
//
116+
// It MUST run before linkCrossRepo (which matches routes by Name). It re-resolves
117+
// from the preserved grpc_service_short prop each run, so it is idempotent across
118+
// appends. A client route whose service has no proto in the snapshot (or an
119+
// ambiguous short name) is left provisional.
120+
func (e *Engine) resolvePyGRPCClientRoutes() {
121+
// Proto index from server routes: short service name → fq, plus per-fq method
122+
// sets and an ambiguity guard (two packages sharing a short service name).
123+
fqOf := map[string]string{}
124+
ambiguous := map[string]bool{}
125+
methodsOf := map[string]map[string]bool{}
126+
for _, r := range e.store.ByKind(facts.KindRoute) {
127+
if valProp(r, "type") != "grpc" || valProp(r, "role") != "server" {
128+
continue
129+
}
130+
fq := valProp(r, "rpc_service")
131+
if fq == "" {
132+
continue
133+
}
134+
short := lastDotSegment(fq)
135+
if prev, ok := fqOf[short]; ok && prev != fq {
136+
ambiguous[short] = true
137+
} else {
138+
fqOf[short] = fq
139+
}
140+
if methodsOf[fq] == nil {
141+
methodsOf[fq] = map[string]bool{}
142+
}
143+
if m := valProp(r, "rpc_method"); m != "" {
144+
methodsOf[fq][m] = true
145+
}
146+
}
147+
148+
// Collect + resolve client routes, then remove-and-re-add so the store's name
149+
// index stays consistent (in-place Name mutation would desync byName).
150+
var replaced []facts.Fact
151+
found := false
152+
for _, r := range e.store.ByKind(facts.KindRoute) {
153+
if valProp(r, "source") != "python-grpc-client" {
154+
continue
155+
}
156+
found = true
157+
short := valProp(r, "grpc_service_short")
158+
method := valProp(r, "rpc_method")
159+
fq := fqOf[short]
160+
if short != "" && method != "" && fq != "" && !ambiguous[short] && methodsOf[fq][method] {
161+
r.Props = cloneProps(r.Props)
162+
r.Name = "/" + fq + "/" + method
163+
r.Props["rpc_service"] = fq
164+
}
165+
replaced = append(replaced, r)
166+
}
167+
if !found {
168+
return
169+
}
170+
e.store.RemoveWhere(func(f facts.Fact) bool {
171+
return f.Kind == facts.KindRoute && valProp(f, "source") == "python-grpc-client"
172+
})
173+
e.store.Add(replaced...)
174+
}
175+
176+
// valProp reads a string prop from a value Fact, tolerating a nil Props map.
177+
func valProp(f facts.Fact, key string) string {
178+
if f.Props == nil {
179+
return ""
180+
}
181+
v, _ := f.Props[key].(string)
182+
return v
183+
}
184+
185+
// cloneProps returns a shallow copy of a props map so a rewritten fact does not
186+
// mutate the shared original.
187+
func cloneProps(p map[string]any) map[string]any {
188+
out := make(map[string]any, len(p)+1)
189+
for k, v := range p {
190+
out[k] = v
191+
}
192+
return out
193+
}
194+
195+
// implShortName returns the gRPC service short name a symbol implements — a Go
196+
// struct embedding Unimplemented<Service>Server, or a Python class subclassing a
197+
// generated <Service>Servicer base — or "" if it implements no such type.
104198
func implShortName(s facts.Fact) string {
105199
for _, r := range s.Relations {
106200
if r.Kind != facts.RelImplements {
@@ -109,6 +203,9 @@ func implShortName(s facts.Fact) string {
109203
if m := unimplementedEmbed.FindStringSubmatch(r.Target); m != nil {
110204
return m[1]
111205
}
206+
if m := servicerBase.FindStringSubmatch(r.Target); m != nil {
207+
return m[1]
208+
}
112209
}
113210
return ""
114211
}

internal/engine/grpcbind_test.go

Lines changed: 94 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -65,6 +65,100 @@ func TestBindGRPCHandlers_BindsViaEmbedConvention(t *testing.T) {
6565
}
6666
}
6767

68+
// pyGRPCClientRoute builds a provisional (short-name) Python gRPC client route like
69+
// the Python extractor emits before the engine resolves it to the FQ wire path.
70+
func pyGRPCClientRoute(short, method string) facts.Fact {
71+
return facts.Fact{
72+
Kind: facts.KindRoute,
73+
Name: "/" + short + "/" + method,
74+
Props: map[string]any{
75+
"role": "client", "method": "POST", "framework": "grpc",
76+
"language": "python", "source": "python-grpc-client", "type": "grpc",
77+
"rpc_method": method, "grpc_service_short": short,
78+
},
79+
}
80+
}
81+
82+
// TestBindGRPCHandlers_BindsPythonServicer: a Python servicer class subclassing a
83+
// generated <Service>Servicer base binds the proto route to its handler method,
84+
// exactly as the Go embed convention does.
85+
func TestBindGRPCHandlers_BindsPythonServicer(t *testing.T) {
86+
eng, _ := New(config.Default())
87+
eng.Store().Add(
88+
grpcServerRoute("vosk.stt.v1.SttService", "StreamingRecognize"),
89+
facts.Fact{
90+
Kind: facts.KindSymbol, Name: "grpc/stt_server.SttServiceServicer",
91+
Props: map[string]any{"symbol_kind": facts.SymbolClass},
92+
Relations: []facts.Relation{
93+
{Kind: facts.RelImplements, Target: "stt_service_pb2_grpc.SttServiceServicer"},
94+
},
95+
},
96+
facts.Fact{
97+
Kind: facts.KindSymbol, Name: "grpc/stt_server.SttServiceServicer.StreamingRecognize",
98+
Props: map[string]any{"symbol_kind": facts.SymbolMethod},
99+
},
100+
)
101+
102+
eng.bindGRPCHandlers()
103+
104+
routes, _ := eng.Store().QueryAdvanced(facts.QueryOpts{Kind: facts.KindRoute, Name: "/vosk.stt.v1.SttService/StreamingRecognize"})
105+
if len(routes) != 1 {
106+
t.Fatalf("route lookup = %d, want 1", len(routes))
107+
}
108+
target, ok := routeHandledBy(routes[0])
109+
want := "grpc/stt_server.SttServiceServicer.StreamingRecognize"
110+
if !ok || target != want {
111+
t.Errorf("handled_by = %q (ok=%v), want %q", target, ok, want)
112+
}
113+
if routes[0].Props["handler"] != want {
114+
t.Errorf("handler prop = %v, want %q", routes[0].Props["handler"], want)
115+
}
116+
}
117+
118+
// TestResolvePyGRPCClientRoutes_RewritesToFQ: a provisional short client route is
119+
// rewritten to the fully-qualified wire path (matching the proto server route) and
120+
// the short route no longer exists.
121+
func TestResolvePyGRPCClientRoutes_RewritesToFQ(t *testing.T) {
122+
eng, _ := New(config.Default())
123+
eng.Store().Add(
124+
grpcServerRoute("vosk.stt.v1.SttService", "StreamingRecognize"),
125+
pyGRPCClientRoute("SttService", "StreamingRecognize"),
126+
)
127+
128+
eng.resolvePyGRPCClientRoutes()
129+
130+
if got, _ := eng.Store().QueryAdvanced(facts.QueryOpts{Kind: facts.KindRoute, Name: "/SttService/StreamingRecognize"}); len(got) != 0 {
131+
t.Errorf("short client route still present: %d", len(got))
132+
}
133+
routes, _ := eng.Store().QueryAdvanced(facts.QueryOpts{Kind: facts.KindRoute, Name: "/vosk.stt.v1.SttService/StreamingRecognize"})
134+
foundClient := false
135+
for _, r := range routes {
136+
if r.Props["source"] == "python-grpc-client" {
137+
foundClient = true
138+
if r.Props["rpc_service"] != "vosk.stt.v1.SttService" {
139+
t.Errorf("rpc_service = %v, want vosk.stt.v1.SttService", r.Props["rpc_service"])
140+
}
141+
}
142+
}
143+
if !foundClient {
144+
t.Error("client route was not rewritten to the fully-qualified name")
145+
}
146+
}
147+
148+
// TestResolvePyGRPCClientRoutes_UnresolvedLeftShort: with no proto in the snapshot,
149+
// the client route stays provisional (short) rather than being dropped or corrupted.
150+
func TestResolvePyGRPCClientRoutes_UnresolvedLeftShort(t *testing.T) {
151+
eng, _ := New(config.Default())
152+
eng.Store().Add(pyGRPCClientRoute("SttService", "StreamingRecognize"))
153+
154+
eng.resolvePyGRPCClientRoutes()
155+
156+
routes, _ := eng.Store().QueryAdvanced(facts.QueryOpts{Kind: facts.KindRoute, Name: "/SttService/StreamingRecognize"})
157+
if len(routes) != 1 {
158+
t.Errorf("unresolved client route should be left short, got %d routes", len(routes))
159+
}
160+
}
161+
68162
func TestBindGRPCHandlers_NoEmbedNoBind(t *testing.T) {
69163
eng, _ := New(config.Default())
70164
eng.Store().Add(

0 commit comments

Comments
 (0)