Skip to content

Commit 8b53631

Browse files
shivasuryaclaude
shivasurya
and
claude
authored
feat(callgraph): Python Type Inference for Improved Call Resolution (#334)
* feat(callgraph): Add type inference data structures Add foundational data structures for type inference: - TypeInfo: tracks type FQN, confidence, and source - VariableBinding: tracks variable types within scopes - FunctionScope: maintains type environment per function - TypeInferenceEngine: coordinates type inference across codebase Tests: 100% coverage with 13 test cases 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com> * feat(callgraph): Add Python builtin type registry Add comprehensive builtin type registry for Python types: - BuiltinRegistry with all Python builtin types (str, list, dict, set, tuple, int, float, bool, bytes) - Method definitions with return types for each builtin - InferLiteralType for automatic type inference from literals - Support for numeric literals (int, float, scientific notation, hex, octal, binary) - Support for collection literals (list, dict, set, tuple) - Support for string and bytes literals Integrate builtin registry with TypeInferenceEngine Tests: 100% coverage with 16 test cases covering all types and methods 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com> * feat(callgraph): Add variable assignment extraction Add AST traversal to extract variable assignments for type inference: - ExtractVariableAssignments traverses Python AST to find assignments - processAssignment extracts type from RHS expression - inferTypeFromExpression handles literal type inference - Supports string, numeric, collection, bool, and None literals - Tracks variable bindings per function scope - Records source locations for each assignment Integration with type inference engine: - Populates function scopes with variable bindings - Handles nested function scopes - Supports variable reassignment (last wins) Tests: 83-93% coverage with 11 test cases covering all literal types 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com> * feat(callgraph): Integrate type inference into call resolution Integrate type inference engine with call graph resolution: - Initialize TypeInferenceEngine in BuildCallGraph - Extract variable assignments during graph building - Use type inference to resolve variable.method() calls - Resolve builtin method calls (str.upper, list.append, dict.keys, etc.) - Fallback to legacy resolution for backward compatibility Type-aware resolution logic: - Check variable bindings in function scope - Resolve builtin methods via BuiltinRegistry - Support user-defined type methods - Maintains 100% backward compatibility with existing tests Integration tests: 5 test cases covering string, list, dict methods and edge cases - TestTypeInference_StringMethods: data.upper() resolution - TestTypeInference_ListMethods: numbers.append(), numbers.count() - TestTypeInference_DictMethods: config.keys(), config.values() - TestTypeInference_MultipleVariables: mixed type resolution - TestTypeInference_WithoutTypeInfo: graceful fallback Backward compatibility: - resolveCallTargetLegacy for tests without type engine - All existing tests pass unchanged 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com> * fix(callgraph): Fix linting issues in type inference code - Convert if-else chain to switch statement in isNumericLiteral - Add period to TODO comment - Add nolint directives for disabled test function 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com> --------- Co-authored-by: Claude <noreply@anthropic.com>
1 parent 2a0be94 commit 8b53631

11 files changed

Lines changed: 2918 additions & 20 deletions

‎sourcecode-parser/graph/callgraph/benchmark_test.go‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -346,12 +346,12 @@ func BenchmarkResolveCallTarget(b *testing.B) {
346346

347347
for i := 0; i < b.N; i++ {
348348
// Test simple attribute access (most common case)
349-
_, _ = resolveCallTarget("utils.process_data", importMap, registry, currentModule, codeGraph)
349+
_, _ = resolveCallTarget("utils.process_data", importMap, registry, currentModule, codeGraph, nil, "")
350350

351351
// Test aliased import
352-
_, _ = resolveCallTarget("helper.format", importMap, registry, currentModule, codeGraph)
352+
_, _ = resolveCallTarget("helper.format", importMap, registry, currentModule, codeGraph, nil, "")
353353

354354
// Test fully qualified name
355-
_, _ = resolveCallTarget("myapp.utils.validate", importMap, registry, currentModule, codeGraph)
355+
_, _ = resolveCallTarget("myapp.utils.validate", importMap, registry, currentModule, codeGraph, nil, "")
356356
}
357357
}

‎sourcecode-parser/graph/callgraph/builder.go‎

Lines changed: 122 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -138,6 +138,10 @@ func BuildCallGraph(codeGraph *graph.CodeGraph, registry *ModuleRegistry, projec
138138
// This avoids re-parsing imports from the same file multiple times
139139
importCache := NewImportMapCache()
140140

141+
// Initialize type inference engine
142+
typeEngine := NewTypeInferenceEngine(registry)
143+
typeEngine.Builtins = NewBuiltinRegistry()
144+
141145
// First, index all function definitions from the code graph
142146
// This builds the Functions map for quick lookup
143147
indexFunctions(codeGraph, callGraph, registry)
@@ -163,6 +167,9 @@ func BuildCallGraph(codeGraph *graph.CodeGraph, registry *ModuleRegistry, projec
163167
continue
164168
}
165169

170+
// Extract variable assignments for type inference
171+
_ = ExtractVariableAssignments(filePath, sourceCode, typeEngine, registry, typeEngine.Builtins)
172+
166173
// Extract all call sites from this file
167174
callSites, err := ExtractCallSites(filePath, sourceCode, importMap)
168175
if err != nil {
@@ -183,7 +190,7 @@ func BuildCallGraph(codeGraph *graph.CodeGraph, registry *ModuleRegistry, projec
183190
}
184191

185192
// Resolve the call target to a fully qualified name
186-
targetFQN, resolved := resolveCallTarget(callSite.Target, importMap, registry, modulePath, codeGraph)
193+
targetFQN, resolved := resolveCallTarget(callSite.Target, importMap, registry, modulePath, codeGraph, typeEngine, callerFQN)
187194

188195
// Update call site with resolution information
189196
callSite.TargetFQN = targetFQN
@@ -423,7 +430,11 @@ func categorizeResolutionFailure(target, targetFQN string) string {
423430
return "unknown"
424431
}
425432

426-
func resolveCallTarget(target string, importMap *ImportMap, registry *ModuleRegistry, currentModule string, codeGraph *graph.CodeGraph) (string, bool) {
433+
func resolveCallTarget(target string, importMap *ImportMap, registry *ModuleRegistry, currentModule string, codeGraph *graph.CodeGraph, typeEngine *TypeInferenceEngine, callerFQN string) (string, bool) {
434+
// Backward compatibility: if typeEngine or callerFQN not provided, skip type inference
435+
if typeEngine == nil || callerFQN == "" {
436+
return resolveCallTargetLegacy(target, importMap, registry, currentModule, codeGraph)
437+
}
427438
// Handle self.method() calls - resolve to current module
428439
if strings.HasPrefix(target, "self.") {
429440
methodName := strings.TrimPrefix(target, "self.")
@@ -472,6 +483,30 @@ func resolveCallTarget(target string, importMap *ImportMap, registry *ModuleRegi
472483
base := parts[0]
473484
rest := parts[1]
474485

486+
// Try type inference for variable.method() calls
487+
if typeEngine != nil && callerFQN != "" {
488+
scope := typeEngine.GetScope(callerFQN)
489+
if scope != nil {
490+
// Check if base is a known variable
491+
if binding, exists := scope.Variables[base]; exists && binding.Type != nil {
492+
varTypeFQN := binding.Type.TypeFQN
493+
// Check if it's a builtin type
494+
if typeEngine.Builtins != nil {
495+
method := typeEngine.Builtins.GetMethod(varTypeFQN, rest)
496+
if method != nil {
497+
// Resolved to builtin method
498+
return varTypeFQN + "." + rest, true
499+
}
500+
}
501+
// Try to resolve as user-defined type method
502+
fullFQN := varTypeFQN + "." + rest
503+
if validateFQN(fullFQN, registry) {
504+
return fullFQN, true
505+
}
506+
}
507+
}
508+
}
509+
475510
// Try to resolve base through imports
476511
if baseFQN, ok := importMap.Resolve(base); ok {
477512
fullFQN := baseFQN + "." + rest
@@ -538,6 +573,91 @@ func validateFQN(fqn string, registry *ModuleRegistry) bool {
538573
return false
539574
}
540575

576+
// resolveCallTargetLegacy is the old resolution logic without type inference.
577+
// Used for backward compatibility with existing tests.
578+
func resolveCallTargetLegacy(target string, importMap *ImportMap, registry *ModuleRegistry, currentModule string, codeGraph *graph.CodeGraph) (string, bool) {
579+
// Handle self.method() calls - resolve to current module
580+
if strings.HasPrefix(target, "self.") {
581+
methodName := strings.TrimPrefix(target, "self.")
582+
// Resolve to module.method
583+
moduleFQN := currentModule + "." + methodName
584+
// Validate exists
585+
if validateFQN(moduleFQN, registry) {
586+
return moduleFQN, true
587+
}
588+
// Return unresolved but with module prefix
589+
return moduleFQN, false
590+
}
591+
592+
// Handle simple names (no dots)
593+
if !strings.Contains(target, ".") {
594+
// Check if it's a Python built-in
595+
if pythonBuiltins[target] {
596+
// Return as builtins.function for pattern matching
597+
return "builtins." + target, true
598+
}
599+
600+
// Try to resolve through imports
601+
if fqn, ok := importMap.Resolve(target); ok {
602+
// Found in imports - return the FQN
603+
// Check if it's a known framework
604+
if isKnown, _ := IsKnownFramework(fqn); isKnown {
605+
return fqn, true
606+
}
607+
// Validate if it exists in registry
608+
resolved := validateFQN(fqn, registry)
609+
return fqn, resolved
610+
}
611+
612+
// Not in imports - might be in same module
613+
sameLevelFQN := currentModule + "." + target
614+
if validateFQN(sameLevelFQN, registry) {
615+
return sameLevelFQN, true
616+
}
617+
618+
// Can't resolve - return as-is
619+
return target, false
620+
}
621+
622+
// Handle qualified names (with dots)
623+
parts := strings.SplitN(target, ".", 2)
624+
base := parts[0]
625+
rest := parts[1]
626+
627+
// Try to resolve base through imports
628+
if baseFQN, ok := importMap.Resolve(base); ok {
629+
fullFQN := baseFQN + "." + rest
630+
// Check if it's a known framework
631+
if isKnown, _ := IsKnownFramework(fullFQN); isKnown {
632+
return fullFQN, true
633+
}
634+
// Check if it's an ORM pattern (before validateFQN, since ORM methods don't exist in source)
635+
if ormFQN, resolved := ResolveORMCall(target, currentModule, registry, codeGraph); resolved {
636+
return ormFQN, true
637+
}
638+
if validateFQN(fullFQN, registry) {
639+
return fullFQN, true
640+
}
641+
return fullFQN, false
642+
}
643+
644+
// Base not in imports - might be module-level access
645+
// Try current module
646+
fullFQN := currentModule + "." + target
647+
if validateFQN(fullFQN, registry) {
648+
return fullFQN, true
649+
}
650+
651+
// Before giving up, check if it's an ORM pattern (Django, SQLAlchemy, etc.)
652+
// ORM methods are dynamically generated at runtime and won't be in source
653+
if ormFQN, resolved := ResolveORMCall(target, currentModule, registry, codeGraph); resolved {
654+
return ormFQN, true
655+
}
656+
657+
// Can't resolve - return as-is
658+
return target, false
659+
}
660+
541661
// readFileBytes reads a file and returns its contents as a byte slice.
542662
// Helper function for reading source code.
543663
func readFileBytes(filePath string) ([]byte, error) {

‎sourcecode-parser/graph/callgraph/builder_framework_test.go‎

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -68,27 +68,27 @@ def test_stdlib():
6868
codeGraph := &graph.CodeGraph{Nodes: make(map[string]*graph.Node)}
6969

7070
// Test Django models resolution
71-
targetFQN, resolved := resolveCallTarget("models.User", importMap, registry, modulePath, codeGraph)
71+
targetFQN, resolved := resolveCallTarget("models.User", importMap, registry, modulePath, codeGraph, nil, "")
7272
assert.True(t, resolved, "Django models.User should be resolved")
7373
assert.Equal(t, "django.db.models.User", targetFQN)
7474

7575
// Test REST framework resolution
76-
targetFQN, resolved = resolveCallTarget("serializers.ModelSerializer", importMap, registry, modulePath, codeGraph)
76+
targetFQN, resolved = resolveCallTarget("serializers.ModelSerializer", importMap, registry, modulePath, codeGraph, nil, "")
7777
assert.True(t, resolved, "REST framework serializers should be resolved")
7878
assert.Equal(t, "rest_framework.serializers.ModelSerializer", targetFQN)
7979

8080
// Test pytest resolution
81-
targetFQN, resolved = resolveCallTarget("pytest.fixture", importMap, registry, modulePath, codeGraph)
81+
targetFQN, resolved = resolveCallTarget("pytest.fixture", importMap, registry, modulePath, codeGraph, nil, "")
8282
assert.True(t, resolved, "pytest.fixture should be resolved")
8383
assert.Equal(t, "pytest.fixture", targetFQN)
8484

8585
// Test json (stdlib) resolution
86-
targetFQN, resolved = resolveCallTarget("json.loads", importMap, registry, modulePath, codeGraph)
86+
targetFQN, resolved = resolveCallTarget("json.loads", importMap, registry, modulePath, codeGraph, nil, "")
8787
assert.True(t, resolved, "json.loads should be resolved")
8888
assert.Equal(t, "json.loads", targetFQN)
8989

9090
// Test logging (stdlib) resolution
91-
targetFQN, resolved = resolveCallTarget("logging.getLogger", importMap, registry, modulePath, codeGraph)
91+
targetFQN, resolved = resolveCallTarget("logging.getLogger", importMap, registry, modulePath, codeGraph, nil, "")
9292
assert.True(t, resolved, "logging.getLogger should be resolved")
9393
assert.Equal(t, "logging.getLogger", targetFQN)
9494
}
@@ -143,11 +143,11 @@ def process():
143143
codeGraph := &graph.CodeGraph{Nodes: make(map[string]*graph.Node)}
144144

145145
// Test local function resolution (should resolve to local module)
146-
targetFQN, resolved := resolveCallTarget("sanitize", importMap, registry, modulePath, codeGraph)
146+
targetFQN, resolved := resolveCallTarget("sanitize", importMap, registry, modulePath, codeGraph, nil, "")
147147
assert.True(t, resolved, "Local function sanitize should be resolved")
148148
assert.Contains(t, targetFQN, "utils.sanitize")
149149

150-
targetFQN, resolved = resolveCallTarget("validate", importMap, registry, modulePath, codeGraph)
150+
targetFQN, resolved = resolveCallTarget("validate", importMap, registry, modulePath, codeGraph, nil, "")
151151
assert.True(t, resolved, "Local function validate should be resolved")
152152
assert.Contains(t, targetFQN, "utils.validate")
153153
}
@@ -197,7 +197,7 @@ def process():
197197
codeGraph := &graph.CodeGraph{Nodes: make(map[string]*graph.Node)}
198198

199199
// Test that local json takes precedence over stdlib
200-
targetFQN, resolved := resolveCallTarget("loads", importMap, registry, modulePath, codeGraph)
200+
targetFQN, resolved := resolveCallTarget("loads", importMap, registry, modulePath, codeGraph, nil, "")
201201
assert.True(t, resolved, "Local json.loads should be resolved")
202202
// When there's a local module that shadows stdlib, it resolves to local
203203
// The FQN will be json.loads but from the local module, not stdlib
@@ -257,12 +257,12 @@ def process():
257257
codeGraph := &graph.CodeGraph{Nodes: make(map[string]*graph.Node)}
258258

259259
// Test local function resolution
260-
targetFQN, resolved := resolveCallTarget("helper", importMap, registry, modulePath, codeGraph)
260+
targetFQN, resolved := resolveCallTarget("helper", importMap, registry, modulePath, codeGraph, nil, "")
261261
assert.True(t, resolved, "Local helper should be resolved")
262262
assert.Contains(t, targetFQN, "utils.helper")
263263

264264
// Test framework resolution
265-
targetFQN, resolved = resolveCallTarget("json.loads", importMap, registry, modulePath, codeGraph)
265+
targetFQN, resolved = resolveCallTarget("json.loads", importMap, registry, modulePath, codeGraph, nil, "")
266266
assert.True(t, resolved, "json.loads should be resolved as framework")
267267
assert.Equal(t, "json.loads", targetFQN)
268268
}

‎sourcecode-parser/graph/callgraph/builder_test.go‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@ func TestResolveCallTarget_SimpleImportedFunction(t *testing.T) {
2323
importMap.AddImport("sanitize", "myapp.utils.sanitize")
2424

2525
codeGraph := &graph.CodeGraph{Nodes: make(map[string]*graph.Node)}
26-
fqn, resolved := resolveCallTarget("sanitize", importMap, registry, "myapp.views", codeGraph)
26+
fqn, resolved := resolveCallTarget("sanitize", importMap, registry, "myapp.views", codeGraph, nil, "")
2727

2828
assert.True(t, resolved)
2929
assert.Equal(t, "myapp.utils.sanitize", fqn)
@@ -42,7 +42,7 @@ func TestResolveCallTarget_QualifiedImport(t *testing.T) {
4242
importMap.AddImport("utils", "myapp.utils")
4343

4444
codeGraph := &graph.CodeGraph{Nodes: make(map[string]*graph.Node)}
45-
fqn, resolved := resolveCallTarget("utils.sanitize", importMap, registry, "myapp.views", codeGraph)
45+
fqn, resolved := resolveCallTarget("utils.sanitize", importMap, registry, "myapp.views", codeGraph, nil, "")
4646

4747
assert.True(t, resolved)
4848
assert.Equal(t, "myapp.utils.sanitize", fqn)
@@ -58,7 +58,7 @@ func TestResolveCallTarget_SameModuleFunction(t *testing.T) {
5858
importMap := NewImportMap("/project/myapp/views.py")
5959

6060
codeGraph := &graph.CodeGraph{Nodes: make(map[string]*graph.Node)}
61-
fqn, resolved := resolveCallTarget("helper", importMap, registry, "myapp.views", codeGraph)
61+
fqn, resolved := resolveCallTarget("helper", importMap, registry, "myapp.views", codeGraph, nil, "")
6262

6363
assert.True(t, resolved)
6464
assert.Equal(t, "myapp.views.helper", fqn)
@@ -74,7 +74,7 @@ func TestResolveCallTarget_UnresolvedMethodCall(t *testing.T) {
7474
importMap := NewImportMap("/project/myapp/views.py")
7575

7676
codeGraph := &graph.CodeGraph{Nodes: make(map[string]*graph.Node)}
77-
fqn, resolved := resolveCallTarget("obj.method", importMap, registry, "myapp.views", codeGraph)
77+
fqn, resolved := resolveCallTarget("obj.method", importMap, registry, "myapp.views", codeGraph, nil, "")
7878

7979
assert.False(t, resolved)
8080
assert.Equal(t, "obj.method", fqn)
@@ -90,7 +90,7 @@ func TestResolveCallTarget_NonExistentFunction(t *testing.T) {
9090
importMap.AddImport("missing", "nonexistent.module.function")
9191

9292
codeGraph := &graph.CodeGraph{Nodes: make(map[string]*graph.Node)}
93-
fqn, resolved := resolveCallTarget("missing", importMap, registry, "myapp.views", codeGraph)
93+
fqn, resolved := resolveCallTarget("missing", importMap, registry, "myapp.views", codeGraph, nil, "")
9494

9595
assert.False(t, resolved)
9696
assert.Equal(t, "nonexistent.module.function", fqn)

0 commit comments

Comments
 (0)