Skip to content

graph call_extractor

Jump to bottom
aakash-anko edited this page May 25, 2026 · 1 revision

call_extractor.py

Extracts function/method call sites from source files using tree-sitter. Supports 14 languages. Produces a list of {caller, callee_name, line} dicts that get resolved against the symbols table in GraphStore._populate_symbol_calls().

Key Concepts

Term Definition Example
AST Abstract Syntax Tree — a tree representation of source code structure, where each node is a language construct (function, class, if-statement, etc.). def add(a, b): return a+b becomes a tree: FunctionDef → [args: a, b] → [body: Return → BinOp(a + b)].
tree-sitter A fast, multi-language parser that builds ASTs. Supports 100+ languages without needing each language's compiler. tree-sitter parses config.py into an AST, then we extract function/class nodes from it.

Module-level constants

CALL_TYPES (Line 9)

A dict mapping language names to their tree-sitter node types that represent function calls.

CALL_TYPES = {
    "python":     ["call"],
    "javascript": ["call_expression"],
    "typescript": ["call_expression"],
    "dart":       ["function_expression_invocation", "method_invocation"],
    "java":       ["method_invocation"],
    "go":         ["call_expression"],
    "rust":       ["call_expression"],
    "ruby":       ["call", "method_call"],
    "c":          ["call_expression"],
    "cpp":        ["call_expression"],
    "csharp":     ["invocation_expression"],
    "php":        ["function_call_expression", "method_call_expression"],
    "kotlin":     ["call_expression"],
    "swift":      ["call_expression"],
}

_FUNCTION_FIELDS (Line 27)

Tuple of tree-sitter field names to check when extracting the callee from a call node:

_FUNCTION_FIELDS = ("function", "name", "method")

_NAME_FIELDS (Line 29)

Tuple of field names to check when extracting the final name from a member expression:

_NAME_FIELDS = ("property", "attribute", "field", "name")

_MEMBER_TYPES (Line 31)

Frozenset of tree-sitter node types that represent member access across all 14 languages (e.g., obj.method, pkg.Func, mod::func):

_MEMBER_TYPES = frozenset({
    "attribute", "member_expression", "selector_expression",
    "field_expression", "member_access_expression", "scoped_identifier",
    "qualified_name", "navigation_expression",
})

Function: _extract_callee_name(call_node) -> str | None (Line 42)

Extracts the function/method name being called from a tree-sitter call expression node. Handles plain calls (foo()), method calls (self.foo()), and qualified calls (Mod.func()).

Example

Input: call_node = tree-sitter node for `self.save()`
       call_node.type = "call"
       call_node has child_by_field_name("function") → attribute node for "self.save"

Lines 53–56 — Try each field in _FUNCTION_FIELDS = ("function", "name", "method"):

  • call_node.child_by_field_name("function") → returns an attribute node for self.save 
func_node = <attribute node: "self.save">
func_node.type = "attribute"

Line 58 — func_node is None? No, skip the fallback.

Line 64 — func_node.type in ("identifier", ...) → "attribute" is NOT in this set, skip.

Line 69 — func_node.type in _MEMBER_TYPES → "attribute" IS in _MEMBER_TYPES ✓

Lines 70–73 — Try each field in _NAME_FIELDS = ("property", "attribute", "field", "name"):

  • func_node.child_by_field_name("property") → None
  • func_node.child_by_field_name("attribute") → returns identifier node "save" 
name_child = <identifier node: "save">
name_child.text = b"save"

Line 73 — name_child.text.decode("utf-8") → "save"

Returns: "save"

Example 2: Plain call foo() 

Input: call_node for `foo()`

Line 54 — child_by_field_name("function") → identifier node "foo", type = "identifier"

Line 65 — func_node.type in ("identifier", ...) → Yes ✓

Line 67 — func_node.text.decode("utf-8") → "foo"

Returns: "foo"

Example 3: No function field found, bare call

Input: call_node with no "function"/"name"/"method" fields
       but has a child with type "identifier" = "print"

Lines 53–56 — All child_by_field_name calls return None. func_node stays None.

Lines 59–61 — Fallback: iterate call_node.children, find first child with type "identifier":

child.type = "identifier", child.text = b"print"

Returns: "print"

Returns None if no callee name can be extracted.

Function: extract_calls_from_file(file_path, language, identifier_path=None) -> list[dict] (Line 83)

Extracts all call sites from a single source file. Uses iterative DFS over the tree-sitter AST to find call nodes, then resolves each call's scope (which function/class/module contains it).

Example

Input:
  file_path = "/repo/main.py"
  language = "python"
  identifier_path = "main.py"

File contents of main.py:
  1: import config
  2:
  3: def run():
  4:     config.load()
  5:     save()
  6:
  7: print("hello")

Line 112 — parser = get_parser_for_language("python") → a tree-sitter Parser instance

Line 116 — node_config = NODE_TYPES.get("python"):

node_config = {"function": ["function_definition"], "class": ["class_definition"], "name_field": "name"}

Line 117 — call_types = CALL_TYPES.get("python"):

call_types = ["call"]

Lines 123–124 — Read the file:

source = b"import config\n\ndef run():\n    config.load()\n    save()\n\nprint(\"hello\")\n"

Line 127 — parser.parse(source) → a tree-sitter Tree

Line 129 — call_type_set = {"call"}

Line 130 — function_types = {"function_definition"}

Line 131 — class_types = {"class_definition"}

Line 132 — all_def_types = {"function_definition", "class_definition"}

Lines 136–137 — Initialize:

module_scope = "main.py:<module>"
stack = [(root_node, "main.py:<module>")]

DFS iteration — processing nodes:

The DFS walks every node. When it hits a function_definition node, it updates current_scope. When it hits a call node, it extracts the callee.

Encountering def run()::

  • node.type = "function_definition" → in all_def_types
  • extract_name(node, "name") → "run"
  • current_scope = "main.py:run"

Encountering config.load() inside run:

  • node.type = "call" → in call_type_set
  • _extract_callee_name(node) → "load"
  • current_scope = "main.py:run", caller_short = "run"
  • callee = "load" ≠ "run" → not self-recursion
  • line = node.start_point[0] + 1 → 4
  • key = ("main.py:run", "load", 4) → not in seen, add it
  • Append: {"caller": "main.py:run", "callee_name": "load", "line": 4}

Encountering save() inside run:

  • node.type = "call" → in call_type_set
  • _extract_callee_name(node) → "save"
  • line = 5
  • Append: {"caller": "main.py:run", "callee_name": "save", "line": 5}

Encountering print("hello") at module level:

  • current_scope = "main.py:<module>"
  • _extract_callee_name(node) → "print"
  • caller_short = "<module>", callee = "print" ≠ "<module>"
  • line = 7
  • Append: {"caller": "main.py:<module>", "callee_name": "print", "line": 7}

Returns:

[
  {"caller": "main.py:run", "callee_name": "load", "line": 4},
  {"caller": "main.py:run", "callee_name": "save", "line": 5},
  {"caller": "main.py:<module>", "callee_name": "print", "line": 7}
]

Scope tracking rules:

Call location Scope assigned
Inside a function "file_path:function_name"
Inside a class but outside methods "file_path:class_name"
At module level (top of file) "file_path:<module>"

Self-recursion filter (line 155):

Calls where the callee name equals the caller name are skipped. E.g., inside def run(): run(), the run() call is NOT recorded because callee == caller_short.

Deduplication (lines 157–159):

The seen set tracks (caller, callee, line) tuples. If the same call site appears twice in the AST traversal (can happen with nested nodes), the duplicate is skipped.

Function: extract_calls_batch(files: list[dict]) -> list[dict] (Line 166)

Extracts call sites from all files that have tree-sitter support. Wrapper around extract_calls_from_file for batch processing.

Example

Input: files = [
  {"file_path": "main.py", "language": "python", "absolute_path": "/repo/main.py"},
  {"file_path": "style.css", "language": "css", "absolute_path": "/repo/style.css"},
  {"file_path": "utils.py", "language": "python", "absolute_path": "/repo/utils.py"}
]

Lines 177–178 — Initialize:

all_calls = []
parsed = 0
skipped = 0

File 1: main.py

  • Line 181 — language = "python"
  • Line 182 — "python" in CALL_TYPES → True, don't skip
  • Line 186 — read_path = "/repo/main.py"
  • Line 187 — extract_calls_from_file("/repo/main.py", "python", identifier_path="main.py"): Returns e.g. [{"caller": "main.py:run", "callee_name": "load", "line": 4}]
  • Line 188 — all_calls.extend(...) → all_calls has 1 item
  • Line 189 — parsed = 1

File 2: style.css

  • Line 181 — language = "css"
  • Line 182 — "css" in CALL_TYPES → False
  • Line 183 — skipped = 1, continue

File 3: utils.py

  • Line 181 — language = "python"
  • Line 186 — extract_calls_from_file("/repo/utils.py", "python", identifier_path="utils.py"): Returns e.g. [{"caller": "utils.py:helper", "callee_name": "format", "line": 10}]
  • Line 188 — all_calls.extend(...) → all_calls has 2 items
  • Line 189 — parsed = 2

Line 191 — Logs: [call_extractor] Extracted 2 call sites from 2 files (1 skipped — no grammar)

Returns:

[
  {"caller": "main.py:run", "callee_name": "load", "line": 4},
  {"caller": "utils.py:helper", "callee_name": "format", "line": 10}
]

Summary

Function Line Purpose
_extract_callee_name 42 Extract callee function name from a call AST node
extract_calls_from_file 83 Extract all call sites from one source file via DFS
extract_calls_batch 166 Batch wrapper: extract calls from all supported files

Clone this wiki locally