Documentation
¶
Overview ¶
Package extraction provides AST-based code extraction utilities for Python source code.
This package uses tree-sitter to extract program statements from Python source files, converting AST nodes into structured Statement objects for analysis.
Example:
statements, err := extraction.ExtractStatements(sourceCode, "myFunction")
if err != nil {
log.Fatal(err)
}
for _, stmt := range statements {
fmt.Printf("Statement type: %s\n", stmt.Type)
}
Index ¶
- func ExtractCStatements(filePath string, sourceCode []byte, functionNode *sitter.Node) ([]*core.Statement, error)
- func ExtractClassAttributes(filePath string, sourceCode []byte, modulePath string, ...) error
- func ExtractCppStatements(filePath string, sourceCode []byte, functionNode *sitter.Node) ([]*core.Statement, error)
- func ExtractFirstReturnType(typeStr string) string
- func ExtractGoReturnTypes(callGraph *core.CallGraph, registry *core.GoModuleRegistry, ...) error
- func ExtractGoReturnTypesWithProgress(callGraph *core.CallGraph, registry *core.GoModuleRegistry, ...) error
- func ExtractGoStatements(filePath string, sourceCode []byte, functionNode *sitter.Node) ([]*core.Statement, error)
- func ExtractGoVarDeclFromNode(node *sitter.Node, sourceCode []byte) []*core.Statement
- func ExtractGoVariableAssignments(filePath string, sourceCode []byte, ...) error
- func ExtractStatements(filePath string, sourceCode []byte, functionNode *sitter.Node) ([]*core.Statement, error)
- func ExtractVariableAssignments(filePath string, sourceCode []byte, typeEngine *resolution.TypeInferenceEngine, ...) error
- func IsBuiltinType(typeStr string) bool
- func ParseGoFile(sourceCode []byte) (*sitter.Tree, error)
- func ParseGoTypeString(typeStr string, registry *core.GoModuleRegistry, filePath string) (*core.TypeInfo, error)
- func ParsePythonFile(sourceCode []byte) (*sitter.Tree, error)
- func StripPointerPrefix(typeStr string) string
- type AttributeAssignment
- type FunctionJob
Constants ¶
This section is empty.
Variables ¶
This section is empty.
Functions ¶
func ExtractCStatements ¶
func ExtractCStatements(filePath string, sourceCode []byte, functionNode *sitter.Node) ([]*core.Statement, error)
ExtractCStatements walks a C function body and produces one *core.Statement per recognised top-level construct (declaration, expression, return, if/for/while/do/switch). The result feeds the CFG builder (PR-10) and the future variable-dependency graph.
Forward declarations and prototypes (no body) yield (nil, nil) — the caller can iterate without nil checks.
The function is a thin wrapper around the shared clikeExtractor; C and C++ share every dispatcher except for the keyword filter and a handful of C++-only AST nodes (`throw_statement`, `try_statement`, `for_range_loop`).
func ExtractClassAttributes ¶
func ExtractClassAttributes( filePath string, sourceCode []byte, modulePath string, typeEngine *resolution.TypeInferenceEngine, attrRegistry *registry.AttributeRegistry, ) error
ExtractClassAttributes extracts all class attributes from a Python file This is Pass 1 & 2 of the attribute extraction algorithm:
Pass 1: Extract class metadata (FQN, methods, file path) Pass 2: Extract attribute assignments (self.attr = value)
Algorithm:
- Parse file with tree-sitter
- Find all class definitions
- For each class: a. Create ClassAttributes entry b. Collect method names c. Scan for self.attr assignments d. Infer types using 6 strategies
Parameters:
- filePath: absolute path to Python file
- sourceCode: file contents
- modulePath: fully qualified module path (e.g., "myapp.models")
- typeEngine: type inference engine with return types and variables
- registry: attribute registry to populate
Returns:
- error if parsing fails
func ExtractCppStatements ¶
func ExtractCppStatements(filePath string, sourceCode []byte, functionNode *sitter.Node) ([]*core.Statement, error)
ExtractCppStatements walks a C++ function body and produces one *core.Statement per recognised construct.
The C and C++ extractors share every dispatcher; the C++ wrapper adds three extra node types via the `extraNodeHandler` hook:
- throw_statement → StatementTypeRaise (with optional CallTarget for `throw std::runtime_error("...")`)
- try_statement → StatementTypeTry with the body in NestedStatements and each catch clause flattened into ElseBranch.
- for_range_loop → StatementTypeFor capturing the loop variable as Def and the iterable expression as Uses.
The keyword filter is `clike.IsCppKeyword`, which inherits all C keywords and adds `class`, `new`, `this`, `static_cast`, etc. so they never appear in Uses.
func ExtractFirstReturnType ¶
ExtractFirstReturnType extracts the first type from multi-return syntax.
Examples:
- "(string, error)" → "string"
- "(int, bool)" → "int"
- "string" → "string" (no change)
- "(User, error)" → "User"
func ExtractGoReturnTypes ¶
func ExtractGoReturnTypes( callGraph *core.CallGraph, registry *core.GoModuleRegistry, typeEngine *resolution.GoTypeInferenceEngine, ) error
ExtractGoReturnTypes extracts return type information from indexed functions.
This function implements Pass 2a of the Go call graph construction pipeline. It processes functions that were indexed in Pass 1, parsing their return type strings (stored in node.ReturnType) into structured TypeInfo objects.
PERFORMANCE: Uses parallel worker pool to process functions concurrently. Progress is reported to stderr every 500 functions processed.
Algorithm:
- Collect all functions with return types into a job queue
- Spawn worker goroutines (based on CPU count)
- Each worker: a) Reads jobs from queue b) Parses return type using ParseGoTypeString() c) Stores result in typeEngine (thread-safe with mutex)
- Report progress every 500 functions
Thread Safety:
This function is thread-safe because typeEngine.AddReturnType() uses mutexes. Multiple workers can call it concurrently.
Parameters:
- callGraph: The call graph with indexed functions (from Pass 1)
- registry: Go module registry for type resolution (from Phase 1)
- typeEngine: Type inference engine to store results (from PR-13)
- showProgress: If true, prints progress to stderr
Returns:
- error: Currently always returns nil (errors are logged but not propagated)
func ExtractGoReturnTypesWithProgress ¶
func ExtractGoReturnTypesWithProgress( callGraph *core.CallGraph, registry *core.GoModuleRegistry, typeEngine *resolution.GoTypeInferenceEngine, showProgress bool, ) error
ExtractGoReturnTypesWithProgress is the internal implementation with progress control.
func ExtractGoStatements ¶
func ExtractGoStatements(filePath string, sourceCode []byte, functionNode *sitter.Node) ([]*core.Statement, error)
ExtractGoStatements extracts top-level statements from a Go function body. Mirrors ExtractStatements (statements.go:15) for the Go tree-sitter grammar.
This is a FLAT extractor: control flow constructs (if, for, switch, select) are skipped — they are handled by BuildGoCFGFromAST (cfg/builder_go.go).
Signature matches ExtractStatements exactly so GenerateGoTaintSummaries can call both extractors with the same pattern.
func ExtractGoVarDeclFromNode ¶
ExtractGoVarDeclFromNode extracts statements from a top-level var_declaration node. Exported for use by GenerateGoTaintSummaries for package-level variable extraction.
func ExtractGoVariableAssignments ¶
func ExtractGoVariableAssignments( filePath string, sourceCode []byte, typeEngine *resolution.GoTypeInferenceEngine, registry *core.GoModuleRegistry, importMap *core.GoImportMap, callGraph *core.CallGraph, ) error
ExtractGoVariableAssignments extracts variable assignments from a Go file and populates the type inference engine with inferred types.
This function implements Pass 2b of the Go call graph construction pipeline. It processes variable assignments by inferring types from RHS expressions and storing the bindings in function scopes.
Algorithm:
- Parse source code with tree-sitter Go parser
- Track function context during AST traversal
- Find assignment nodes (short_var_declaration, assignment_statement)
- For each assignment: a) Extract LHS variable name(s) b) Infer type from RHS expression: - Function call: Look up return type in engine - Literal: Infer builtin type - Variable ref: Copy type from scope - Struct literal: Extract type name c) Create GoVariableBinding d) Add binding to function scope
RHS Type Inference Patterns:
- GetUser() → Look up GetUser return type
- "Alice" → builtin.string
- 42 → builtin.int
- true → builtin.bool
- user → Copy user's type from scope
- User{} → Extract User type
- &User{} → Extract User type (strip &)
Parameters:
- filePath: Absolute path to the Go source file
- sourceCode: Contents of the file as byte array
- typeEngine: Type inference engine with return types populated (from Pass 2a)
- registry: Go module registry for package resolution
- importMap: Import mappings for resolving qualified type names
Returns:
- error: If parsing fails or other critical errors occur
Example:
engine := resolution.NewGoTypeInferenceEngine(registry) // After Pass 2a (return type extraction) err := ExtractGoVariableAssignments(filePath, sourceCode, engine, registry, importMap) // Now engine.Scopes contains variable bindings for each function
Thread Safety:
This function is thread-safe because typeEngine operations use mutexes. Can be called in parallel for different files.
func ExtractStatements ¶
func ExtractStatements(filePath string, sourceCode []byte, functionNode *sitter.Node) ([]*core.Statement, error)
ExtractStatements extracts all statements from a Python function body. It processes assignments, calls, and returns to build def-use chains. Returns a slice of Statement objects or an error if parsing fails.
func ExtractVariableAssignments ¶
func ExtractVariableAssignments( filePath string, sourceCode []byte, typeEngine *resolution.TypeInferenceEngine, registry *core.ModuleRegistry, builtinRegistry *registry.BuiltinRegistry, importMap *core.ImportMap, ) error
ExtractVariableAssignments extracts variable assignments from a Python file and populates the type inference engine with inferred types.
Algorithm:
- Parse source code with tree-sitter Python parser
- Traverse AST to find assignment statements
- For each assignment: - Extract variable name - Infer type from RHS (literal, function call, or method call) - Create VariableBinding with inferred type - Add binding to function scope
Parameters:
- filePath: absolute path to the Python file
- sourceCode: contents of the file as byte array
- typeEngine: type inference engine to populate
- registry: module registry for resolving module paths
- builtinRegistry: builtin types registry for literal inference
- importMap: import mappings for resolving class instantiations from imports
Returns:
- error: if parsing fails
Note: Class context is tracked during AST traversal by detecting class_definition nodes. This enables building class-qualified FQNs (module.ClassName.methodName) that match the FQNs created during function indexing (Pass 1).
func IsBuiltinType ¶
IsBuiltinType checks if a type string is a Go builtin type.
Examples:
- "int" → true
- "string" → true
- "User" → false
func ParseGoFile ¶
ParseGoFile parses a Go source file using tree-sitter. Mirrors ParsePythonFile in statements.go:525 for Python.
func ParseGoTypeString ¶
func ParseGoTypeString( typeStr string, registry *core.GoModuleRegistry, filePath string, ) (*core.TypeInfo, error)
ParseGoTypeString parses a Go type string into a TypeInfo object.
Handles multiple Go type patterns:
- Builtins: "int", "string", "error" → "builtin.X"
- Pointers: "*User" → strip * → "pkg.User"
- Multi-return: "(string, error)" → extract first → "builtin.string"
- Qualified: "models.User" → resolve package → "github.com/myapp/models.User"
- Same-package: "User" → resolve via registry → "github.com/myapp/handlers.User"
Algorithm:
- Normalize and trim whitespace
- Handle multi-return (extract first)
- Strip pointer prefix
- Check if builtin
- Resolve qualified types (with .)
- Resolve same-package types (via registry)
- Fallback to as-is with lower confidence
Parameters:
- typeStr: Go type string from node.ReturnType (e.g., "*User", "(string, error)")
- registry: Go module registry for package resolution (can be nil)
- filePath: Source file path for same-package resolution (can be empty)
Returns:
- TypeInfo object with resolved FQN
- nil if typeStr is empty
- error if parsing fails critically (currently never returns error)
Examples:
- ParseGoTypeString("int", nil, "") → TypeInfo{TypeFQN: "builtin.int", Confidence: 1.0}
- ParseGoTypeString("*User", registry, "models/user.go") → TypeInfo{TypeFQN: "myapp.models.User", Confidence: 0.95}
- ParseGoTypeString("(string, error)", nil, "") → TypeInfo{TypeFQN: "builtin.string", Confidence: 1.0}
func ParsePythonFile ¶
ParsePythonFile parses a Python source file using tree-sitter. Returns the parsed tree or an error.
func StripPointerPrefix ¶
StripPointerPrefix removes the leading * from pointer types. Multiple pointer levels are stripped to the base type.
Examples:
- "*User" → "User"
- "**Config" → "Config"
- "User" → "User" (no change)
Types ¶
type AttributeAssignment ¶
type AttributeAssignment struct {
AttributeName string // Name of the attribute (e.g., "value", "user")
RightSide *sitter.Node // AST node of the right-hand side expression
Node *sitter.Node // Full assignment node
}
AttributeAssignment represents a self.attr = value assignment.
type FunctionJob ¶
FunctionJob represents a function to process for return type extraction.