extraction

package
v0.0.0-...-460d0d3 Latest Latest
Warning

This package is not in the latest version of its module.

Go to latest
Published: May 23, 2026 License: Apache-2.0 Imports: 18 Imported by: 0

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

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:

  1. Parse file with tree-sitter
  2. Find all class definitions
  3. 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

func ExtractFirstReturnType(typeStr string) string

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:

  1. Collect all functions with return types into a job queue
  2. Spawn worker goroutines (based on CPU count)
  3. Each worker: a) Reads jobs from queue b) Parses return type using ParseGoTypeString() c) Stores result in typeEngine (thread-safe with mutex)
  4. 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

func ExtractGoVarDeclFromNode(node *sitter.Node, sourceCode []byte) []*core.Statement

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:

  1. Parse source code with tree-sitter Go parser
  2. Track function context during AST traversal
  3. Find assignment nodes (short_var_declaration, assignment_statement)
  4. 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:

  1. Parse source code with tree-sitter Python parser
  2. Traverse AST to find assignment statements
  3. 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

func IsBuiltinType(typeStr string) bool

IsBuiltinType checks if a type string is a Go builtin type.

Examples:

  • "int" → true
  • "string" → true
  • "User" → false

func ParseGoFile

func ParseGoFile(sourceCode []byte) (*sitter.Tree, error)

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:

  1. Normalize and trim whitespace
  2. Handle multi-return (extract first)
  3. Strip pointer prefix
  4. Check if builtin
  5. Resolve qualified types (with .)
  6. Resolve same-package types (via registry)
  7. 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

func ParsePythonFile(sourceCode []byte) (*sitter.Tree, error)

ParsePythonFile parses a Python source file using tree-sitter. Returns the parsed tree or an error.

func StripPointerPrefix

func StripPointerPrefix(typeStr string) string

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

type FunctionJob struct {
	FQN        string
	ReturnType string
	File       string
}

FunctionJob represents a function to process for return type extraction.

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL