Skip to content

codenib.graph

Graph-related modules for CodeNib.

Modules:

Name Description
backend_alignment

Backend-alignment checks for CodeGraph-producing indexers.

code_graph
dependency

Dependency / impact analysis over the code call-graph.

hierarchy

Repo-level compound graph model for code visualization.

incremental

LSP-based incremental patching for CodeGraph updates.

layers

Layered graph views over a single :class:CodeGraph.

roi_subgraph
setup

Repository-aware setup diagnostics for the optional symbol graph.

traverse_graph

Classes:

Name Description
CodeGraph

A class to represent and manipulate a code graph using igraph.

GraphPatcher

Orchestrates incremental graph updates using LSP.

LSPClient

Synchronous LSP client that communicates with a language server via stdio.

PatcherBase

Base class for language-specific incremental patchers.

MultiGraphIndex

Layer index for a :class:CodeGraph.

ROISubgraph

Region of Interest (ROI) subgraph extraction for code graphs.

Functions:

Name Description
build_graph_layers

Build a layer index for code_graph.

traverse_tree_structure

Traverse tree structure starting from a root node

CodeGraph

CodeGraph(project_root=None)

A class to represent and manipulate a code graph using igraph. start_line uses 0-based index.

Methods:

Name Description
add_file_node

Add a file node to the graph and set it as the current file and scope.

add_symbol_node

Add a symbol node to the graph.

add_symbol_reference

Add a reference to a symbol.

update_current_scope

Update the current scope to the given symbol with its range.

exit_scopes_by_line

Exit scopes that have ended based on current line number.

add_containment_edge

Add a containment edge from current scope to a symbol.

batch_edges

Defer igraph edge insertion until the block exits, flushing all

build_range_indexes

Rebuild per-file line-range indexes from current graph state.

merge_from

Merge another CodeGraph into this one.

query_range

Query the graph by source-file line range (LSP-aligned).

query_range_by_symbol

Range-query by symbol identity (name).

layers

Build a multi-graph layer index over this graph.

layer

Return one named graph layer view.

add_root_node

Add the root node to the graph

add_directory_node

Add a directory node to the graph

save_graph

Save the graph to a pickle file for fast serialization.

load_graph

Load a graph from a pickle file.

get_graph

Get the igraph Graph object.

unified_index

Readable unified_name (full + bare) -> [canonical names], cached.

display_name

Readable label for a canonical name (unified_name or itself).

resolve_symbol

Map a readable/bare symbol to (canonical_name | None, candidates).

get_node_info_by_name

Get information about a node in the graph.

get_node_info_by_id

Get information about a node in the graph.

get_node_content

Get the content of a node in the graph.

get_neighbors

Get the neighbors of a node in the graph.

get_successors

Get the successors of a node_name (outgoing edges).

get_predecessors

Get the predecessors of a node_name (incoming edges).

print_graph_basic_info

Print basic information about the graph.

visualize_graph

Visualize the code graph with different colors for different node and edge types.

Source code in codenib/graph/code_graph.py
def __init__(self, project_root=None):
    # Create a directed graph
    self.graph = ig.Graph(directed=True)
    self.current_file = None
    self.current_scope = None
    self.project_root = project_root
    self.scope_stack = []  # List of {symbol: [start_line, end_line]} dicts
    # Store line ranges for symbols
    self.symbol_ranges = {}
    # Map symbol names to vertex IDs
    self.name_to_vertex = {}

    # Range indexes — per-file, populated by build_range_indexes() after the
    # graph is fully built or after a patcher batch. Pickled with the graph.
    # _file_nodes[file] -> list of (start_line, end_line, vid), unsorted.
    # _file_edge_anchors[file] -> sorted list of (anchor_line, eid).
    self._file_nodes: Dict[str, List[Tuple[int, int, int]]] = {}
    self._file_edge_anchors: Dict[str, List[Tuple[int, int]]] = {}

    # Aux: unified_name (display, may collide) -> list of identity names.
    # Populated by build_range_indexes(). Pickled with the graph.
    # Consumers disambiguate via file context (this layer just lists candidates).
    self._unified_to_names: Dict[str, List[str]] = {}

    # Edge dedup index: (src_id, tgt_id, type, anchor_file, anchor_line) -> eid.
    # Populated lazily on first _add_edge / invalidated to None when callers
    # delete edges or vertices (igraph compacts eids on delete, so cached
    # eids would point at wrong edges). Lookup via _ensure_edge_index().
    # Not pickled — rebuilt on demand after load.
    self._edge_index: Optional[Dict[Tuple, int]] = None

    # Edge insertion buffer, active only inside a `batch_edges()` block.
    # igraph rebuilds its adjacency index on every `add_edges` call, so
    # inserting one edge at a time is O(E^2). Decoders wrap their build
    # loop in `batch_edges()` to defer insertion into a single add_edges
    # flush (O(E)). When None, `_add_edge` inserts immediately as before.
    self._edge_buffer: Optional[List[Tuple[int, int]]] = None
    self._edge_attr_buffer: Optional[List[Tuple]] = None

add_file_node

add_file_node(file_path)

Add a file node to the graph and set it as the current file and scope.

Parameters:

Name Type Description Default
file_path

Path of the file to add

required
Source code in codenib/graph/code_graph.py
def add_file_node(self, file_path):
    """
    Add a file node to the graph and set it as the current file and scope.

    Args:
        file_path: Path of the file to add
    """
    self.current_file = file_path

    # Add vertex for file
    self._add_vertex(file_path, {"type": NODE_TYPE_FILE})

    self.current_scope = file_path
    # File scope has no range (special case)
    self.scope_stack = [{file_path: None}]

add_symbol_node

add_symbol_node(
    symbol, line, scope_start_line=None, scope_end_line=None, symbol_type=None
)

Add a symbol node to the graph.

Parameters:

Name Type Description Default
symbol

Symbol name

required
line

Line number of the symbol

required
scope_start_line

Start line of the symbol's scope (optional)

None
scope_end_line

End line of the symbol's scope (optional)

None
symbol_type

Type of symbol (NODE_TYPE_CLASS, NODE_TYPE_METHOD, NODE_TYPE_FUNCTION) (optional)

None
Source code in codenib/graph/code_graph.py
def add_symbol_node(
    self, symbol, line, scope_start_line=None, scope_end_line=None, symbol_type=None
):
    """
    Add a symbol node to the graph.

    Args:
        symbol: Symbol name
        line: Line number of the symbol
        scope_start_line: Start line of the symbol's scope (optional)
        scope_end_line: End line of the symbol's scope (optional)
        symbol_type: Type of symbol
            (NODE_TYPE_CLASS, NODE_TYPE_METHOD, NODE_TYPE_FUNCTION) (optional)
    """
    # Use specific symbol type if provided, otherwise default to generic symbol
    node_type = symbol_type if symbol_type else NODE_TYPE_SYMBOL

    if scope_start_line is not None and scope_end_line is not None:
        # Store symbol range
        self.symbol_ranges[symbol] = (scope_start_line, scope_end_line)

        # Add symbol vertex with scope range
        self._add_vertex(
            symbol,
            {
                "type": node_type,
                "file": self.current_file,
                "start_line": scope_start_line,
                "end_line": scope_end_line,
                "selection_line": line,
            },
        )
    else:
        # Add symbol vertex without scope range
        self._add_vertex(
            symbol,
            {
                "type": node_type,
                "file": self.current_file,
                "start_line": line,
                "end_line": line,
                "selection_line": line,
            },
        )

add_symbol_reference

add_symbol_reference(
    symbol,
    module_path=None,
    symbol_type=None,
    anchor_file: str | None = None,
    anchor_line: int | None = None,
)

Add a reference to a symbol.

Parameters:

Name Type Description Default
symbol

Symbol being referenced

required
module_path

Path of the module containing the symbol (optional)

None
symbol_type

Type of symbol (NODE_TYPE_CLASS, NODE_TYPE_METHOD, NODE_TYPE_FUNCTION) (optional)

None
anchor_file str | None

File where the reference site is located (optional)

None
anchor_line int | None

0-based line of the reference site (optional)

None
Source code in codenib/graph/code_graph.py
def add_symbol_reference(
    self,
    symbol,
    module_path=None,
    symbol_type=None,
    anchor_file: Optional[str] = None,
    anchor_line: Optional[int] = None,
):
    """
    Add a reference to a symbol.

    Args:
        symbol: Symbol being referenced
        module_path: Path of the module containing the symbol (optional)
        symbol_type: Type of symbol
            (NODE_TYPE_CLASS, NODE_TYPE_METHOD, NODE_TYPE_FUNCTION) (optional)
        anchor_file: File where the reference site is located (optional)
        anchor_line: 0-based line of the reference site (optional)
    """
    # If the symbol doesn't exist, create it without range info
    if symbol not in self.name_to_vertex:
        file_attr = module_path if module_path else None
        node_type = symbol_type if symbol_type else NODE_TYPE_SYMBOL
        self._add_vertex(symbol, {"type": node_type, "file": file_attr})

    # Add reference edge with anchor info (defaults: current file)
    if anchor_file is None:
        anchor_file = self.current_file
    self._add_edge(
        self.current_scope,
        symbol,
        EDGE_TYPE_REFERENCE,
        anchor_file=anchor_file,
        anchor_line=anchor_line,
    )

update_current_scope

update_current_scope(symbol, start_line=None, end_line=None)

Update the current scope to the given symbol with its range.

Parameters:

Name Type Description Default
symbol

Symbol to set as current scope

required
start_line

Start line of the symbol's scope

None
end_line

End line of the symbol's scope

None
Source code in codenib/graph/code_graph.py
def update_current_scope(self, symbol, start_line=None, end_line=None):
    """
    Update the current scope to the given symbol with its range.

    Args:
        symbol: Symbol to set as current scope
        start_line: Start line of the symbol's scope
        end_line: End line of the symbol's scope
    """
    self.current_scope = symbol
    # Add symbol with its range to scope stack
    scope_range = (
        [start_line, end_line]
        if start_line is not None and end_line is not None
        else None
    )
    self.scope_stack.append({symbol: scope_range})

exit_scopes_by_line

exit_scopes_by_line(current_line)

Exit scopes that have ended based on current line number. File scope (with None range) is never popped.

Parameters:

Name Type Description Default
current_line

Current line number being processed

required
Source code in codenib/graph/code_graph.py
def exit_scopes_by_line(self, current_line):
    """
    Exit scopes that have ended based on current line number.
    File scope (with None range) is never popped.

    Args:
        current_line: Current line number being processed
    """
    # Pop scopes whose range has ended
    while len(self.scope_stack) > 1:  # Keep at least the file scope
        top_scope_dict = self.scope_stack[-1]
        scope_symbol = list(top_scope_dict.keys())[0]
        scope_range = top_scope_dict[scope_symbol]

        # If scope has no range (file scope), don't pop
        if scope_range is None:
            break

        # If current line is beyond scope's end, pop it
        start_line, end_line = scope_range
        if current_line > end_line:
            self.scope_stack.pop()
            # Update current_scope to new top of stack
            if self.scope_stack:
                new_top = self.scope_stack[-1]
                self.current_scope = list(new_top.keys())[0]
            else:
                self.current_scope = self.current_file
        else:
            # Scope is still active
            break

add_containment_edge

add_containment_edge(target_symbol)

Add a containment edge from current scope to a symbol.

CONTAIN edges deliberately carry no anchor — containment is a structural relation, not a call/reference site. (See anchor invariant ii on query_range.)

Parameters:

Name Type Description Default
target_symbol

Symbol being contained.

required
Source code in codenib/graph/code_graph.py
def add_containment_edge(self, target_symbol):
    """
    Add a containment edge from current scope to a symbol.

    CONTAIN edges deliberately carry no anchor — containment is a
    structural relation, not a call/reference site. (See anchor invariant
    ii on `query_range`.)

    Args:
        target_symbol: Symbol being contained.
    """
    # Use the current scope directly (not parent scope)
    self._add_edge(self.current_scope, target_symbol, EDGE_TYPE_CONTAIN)

batch_edges

batch_edges()

Defer igraph edge insertion until the block exits, flushing all buffered edges in a single add_edges call.

igraph rebuilds its internal adjacency index on every add_edges call, so inserting one edge at a time (as bare _add_edge does) is O(E^2) in the number of edges — the dominant cost when decoding a large repo. Wrapping a build loop in this context collapses that to O(E):

with graph.batch_edges():
    # ... many add_symbol_reference / _add_edge calls ...

Edge dedup is unaffected (the dedup index is updated as edges are buffered). Edge ids returned by _add_edge while buffering are provisional and only become real after flush; no caller relies on them mid-build. Outside this block _add_edge inserts immediately, so the incremental patcher and all other callers are unchanged. Nesting is a no-op (the outermost block owns the flush).

Source code in codenib/graph/code_graph.py
@contextmanager
def batch_edges(self):
    """Defer igraph edge insertion until the block exits, flushing all
    buffered edges in a single ``add_edges`` call.

    igraph rebuilds its internal adjacency index on every ``add_edges``
    call, so inserting one edge at a time (as bare ``_add_edge`` does) is
    O(E^2) in the number of edges — the dominant cost when decoding a large
    repo. Wrapping a build loop in this context collapses that to O(E):

        with graph.batch_edges():
            # ... many add_symbol_reference / _add_edge calls ...

    Edge dedup is unaffected (the dedup index is updated as edges are
    buffered). Edge *ids* returned by ``_add_edge`` while buffering are
    provisional and only become real after flush; no caller relies on them
    mid-build. Outside this block ``_add_edge`` inserts immediately, so the
    incremental patcher and all other callers are unchanged. Nesting is a
    no-op (the outermost block owns the flush)."""
    if self._edge_buffer is not None:
        # Already batching — let the outer block own the buffer/flush.
        yield
        return
    if self._edge_index is None:
        self._rebuild_edge_index()
    self._edge_buffer = []
    self._edge_attr_buffer = []
    try:
        yield
    finally:
        self._flush_edge_buffer()

build_range_indexes

build_range_indexes() -> None

Rebuild per-file line-range indexes from current graph state.

Call this after the graph is fully constructed by a decoder, and again after any incremental patcher batch. The two indexes are pickled with the graph by save_graph, so subsequent load_graph calls do not need to rebuild.

Cost: O(V + E) — one linear pass over vertices and edges.

Source code in codenib/graph/code_graph.py
def build_range_indexes(self) -> None:
    """Rebuild per-file line-range indexes from current graph state.

    Call this after the graph is fully constructed by a decoder, and
    again after any incremental patcher batch. The two indexes are
    pickled with the graph by `save_graph`, so subsequent `load_graph`
    calls do not need to rebuild.

    Cost: O(V + E) — one linear pass over vertices and edges.
    """
    nodes: Dict[str, List[Tuple[int, int, int]]] = defaultdict(list)
    for vid in range(self.graph.vcount()):
        v = self.graph.vs[vid]
        attrs = v.attributes()
        f = attrs.get("file")
        s = attrs.get("start_line")
        e = attrs.get("end_line")
        if f and s is not None and e is not None:
            nodes[f].append((s, e, vid))
    self._file_nodes = dict(nodes)

    edges: Dict[str, List[Tuple[int, int]]] = defaultdict(list)
    for eid in range(self.graph.ecount()):
        edge = self.graph.es[eid]
        attrs = edge.attributes()
        f = attrs.get("anchor_file")
        line = attrs.get("anchor_line")
        if f and line is not None:
            edges[f].append((line, eid))
    for f in edges:
        edges[f].sort()
    self._file_edge_anchors = dict(edges)

    unified: Dict[str, List[str]] = defaultdict(list)
    for vid in range(self.graph.vcount()):
        v = self.graph.vs[vid]
        uname = v.attributes().get("unified_name")
        if uname:
            unified[uname].append(v["name"])
    self._unified_to_names = dict(unified)
    self._unified_index_cache = None

merge_from

merge_from(other: CodeGraph) -> None

Merge another CodeGraph into this one.

Vertices are keyed by their canonical name attribute. Existing vertices receive any attributes from other; missing vertices are inserted. Edges are copied through _add_edge so structural edge deduplication and anchored reference multi-edges keep the same semantics as decoder-built graphs.

Source code in codenib/graph/code_graph.py
def merge_from(self, other: "CodeGraph") -> None:
    """Merge another ``CodeGraph`` into this one.

    Vertices are keyed by their canonical ``name`` attribute. Existing
    vertices receive any attributes from ``other``; missing vertices are
    inserted. Edges are copied through ``_add_edge`` so structural edge
    deduplication and anchored reference multi-edges keep the same
    semantics as decoder-built graphs.
    """
    for vertex in other.graph.vs:
        attrs = vertex.attributes()
        name = attrs.get("name")
        if not name:
            continue
        copied = {key: value for key, value in attrs.items() if key != "name"}
        self._add_vertex(name, copied)

    with self.batch_edges():
        for edge in other.graph.es:
            attrs = edge.attributes()
            edge_type = attrs.get("type")
            if edge_type is None:
                continue
            source_name = other.graph.vs[edge.source]["name"]
            target_name = other.graph.vs[edge.target]["name"]
            self._add_edge(
                source_name,
                target_name,
                edge_type,
                anchor_file=attrs.get("anchor_file"),
                anchor_line=attrs.get("anchor_line"),
            )

    self.symbol_ranges.update(other.symbol_ranges)
    self.build_range_indexes()

query_range

query_range(
    file: str,
    start_line: int,
    end_line: int,
    kinds: set | None = None,
    depth: int = 1,
    *,
    layer: str | None = None
) -> RangeQueryResult

Query the graph by source-file line range (LSP-aligned).

Parameters:

Name Type Description Default
file str

source file (graph-relative path).

required
start_line int

first line of the inclusive 0-based span.

required
end_line int

last line of the inclusive 0-based span.

required
kinds set | None

edge kinds to include in outgoing AND incoming. Defaults to {EDGE_TYPE_REFERENCE} — CONTAIN edges have no meaningful anchor and are excluded by default. Pass an explicit set (e.g. {EDGE_TYPE_REFERENCE, EDGE_TYPE_CONTAIN}) to opt in.

None
depth int

reserved for future multi-hop expansion. Only depth=1 is supported in v1; raises NotImplementedError otherwise.

1
layer str | None

optional named graph layer (reference, dependency, containment, all, ...). Mutually exclusive with kinds.

None

Returns: RangeQueryResult with typed NodeRef/EdgeRef records.

Anchor invariants

(i) anchor_file == source.file (anchor lives at the call site, never the target). (ii) anchor_* is meaningful only on EDGE_TYPE_REFERENCE; CONTAIN edges leave both fields None. (iii) outgoing ∩ incoming = ∅ for any single query (call site and target are in distinct scopes by definition).

Source code in codenib/graph/code_graph.py
def query_range(
    self,
    file: str,
    start_line: int,
    end_line: int,
    kinds: Optional[set] = None,
    depth: int = 1,
    *,
    layer: Optional[str] = None,
) -> RangeQueryResult:
    """Query the graph by source-file line range (LSP-aligned).

    Args:
        file: source file (graph-relative path).
        start_line: first line of the inclusive 0-based span.
        end_line: last line of the inclusive 0-based span.
        kinds: edge kinds to include in `outgoing` AND `incoming`.
            Defaults to `{EDGE_TYPE_REFERENCE}` — CONTAIN edges have no
            meaningful anchor and are excluded by default. Pass an
            explicit set (e.g. `{EDGE_TYPE_REFERENCE, EDGE_TYPE_CONTAIN}`)
            to opt in.
        depth: reserved for future multi-hop expansion. Only `depth=1`
            is supported in v1; raises NotImplementedError otherwise.
        layer: optional named graph layer (`reference`, `dependency`,
            `containment`, `all`, ...). Mutually exclusive with ``kinds``.

    Returns: RangeQueryResult with typed `NodeRef`/`EdgeRef` records.

    Anchor invariants:
        (i)   anchor_file == source.file (anchor lives at the call site,
              never the target).
        (ii)  anchor_* is meaningful only on EDGE_TYPE_REFERENCE; CONTAIN
              edges leave both fields None.
        (iii) outgoing ∩ incoming = ∅ for any single query (call site
              and target are in distinct scopes by definition).
    """
    if depth != 1:
        raise NotImplementedError(
            f"query_range supports only depth=1 in v1; got depth={depth}"
        )
    if layer is not None:
        if kinds is not None:
            raise ValueError("query_range accepts either kinds or layer, not both")
        layer_kinds = edge_types_for_graph_layer(layer)
        if layer_kinds is None:
            kinds = {edge.attributes().get("type") for edge in self.graph.es}
        else:
            kinds = set(layer_kinds)
    if kinds is None:
        kinds = {EDGE_TYPE_REFERENCE}

    # defined: brute-force overlap on per-file node list.
    nodes = self._file_nodes.get(file, [])
    defined_vids = [vid for s, e, vid in nodes if s <= end_line and e >= start_line]
    defined = [self._build_node_ref(vid) for vid in defined_vids]

    # outgoing: bisect on sorted (anchor_line, eid) list, then filter by kind.
    # Use bisect_left for the upper bound: bisect_right would include
    # `(end_line+1, 0)` (eid==0 entries at the line just past the query),
    # leaking an off-anchor edge into the slice.
    arr = self._file_edge_anchors.get(file, [])
    lo = bisect.bisect_left(arr, (start_line, -1))
    hi = bisect.bisect_left(arr, (end_line + 1, 0))
    outgoing: List[EdgeRef] = []
    for _, eid in arr[lo:hi]:
        edge = self.graph.es[eid]
        if edge["type"] in kinds:
            outgoing.append(self._build_edge_ref(eid))

    # incoming: walk igraph adjacency directly; carry target_vid so the
    # consumer can map the inbound edge back to which defined node it hit.
    # Self-edges (recursion) are emitted by `outgoing` already, so skip
    # them here to keep outgoing ∩ incoming = ∅ (invariant iii).
    incoming: List[EdgeRef] = []
    for vid in defined_vids:
        for eid in self.graph.incident(vid, mode="in"):
            edge = self.graph.es[eid]
            if edge.source == edge.target:
                continue
            if edge["type"] in kinds:
                incoming.append(self._build_edge_ref(eid))

    return RangeQueryResult(defined=defined, outgoing=outgoing, incoming=incoming)

query_range_by_symbol

query_range_by_symbol(
    name: str, kinds: set | None = None, depth: int = 1, *, layer: str | None = None
) -> RangeQueryResult

Range-query by symbol identity (name).

Resolves name -> vid -> (file, start_line, end_line) and forwards to query_range. name is the globally-unique identity attribute (semi-raw SCIP symbol), not the unified_name display.

Returns an empty RangeQueryResult if the symbol is unknown or has no associated file/range (e.g. a SCIP reference-only vertex).

Source code in codenib/graph/code_graph.py
def query_range_by_symbol(
    self,
    name: str,
    kinds: Optional[set] = None,
    depth: int = 1,
    *,
    layer: Optional[str] = None,
) -> RangeQueryResult:
    """Range-query by symbol identity (`name`).

    Resolves `name` -> vid -> `(file, start_line, end_line)` and forwards
    to `query_range`. `name` is the globally-unique identity attribute
    (semi-raw SCIP symbol), not the `unified_name` display.

    Returns an empty RangeQueryResult if the symbol is unknown or has
    no associated file/range (e.g. a SCIP reference-only vertex).
    """
    vid = self.name_to_vertex.get(name)
    if vid is None:
        return RangeQueryResult()
    attrs = self.graph.vs[vid].attributes()
    f = attrs.get("file")
    s = attrs.get("start_line")
    e = attrs.get("end_line")
    if f is None or s is None or e is None:
        return RangeQueryResult()
    return self.query_range(f, s, e, kinds=kinds, depth=depth, layer=layer)

layers

layers(*, use_core: bool = True)

Build a multi-graph layer index over this graph.

Source code in codenib/graph/code_graph.py
def layers(self, *, use_core: bool = True):
    """Build a multi-graph layer index over this graph."""
    from .layers import build_graph_layers

    return build_graph_layers(self, use_core=use_core)

layer

layer(name: str = GRAPH_LAYER_REFERENCE, *, use_core: bool = True)

Return one named graph layer view.

Source code in codenib/graph/code_graph.py
def layer(self, name: str = GRAPH_LAYER_REFERENCE, *, use_core: bool = True):
    """Return one named graph layer view."""
    return self.layers(use_core=use_core).get(name)

add_root_node

add_root_node(project_root)

Add the root node to the graph

Source code in codenib/graph/code_graph.py
def add_root_node(self, project_root):
    """Add the root node to the graph"""
    self._add_vertex(project_root, {"type": "root"})

add_directory_node

add_directory_node(dir_path)

Add a directory node to the graph

Source code in codenib/graph/code_graph.py
def add_directory_node(self, dir_path):
    """Add a directory node to the graph"""
    self._add_vertex(dir_path, {"type": NODE_TYPE_DIRECTORY})

save_graph

save_graph(output_path)

Save the graph to a pickle file for fast serialization.

Parameters:

Name Type Description Default
output_path

Path to save the pickle file

required
Source code in codenib/graph/code_graph.py
def save_graph(self, output_path):
    """
    Save the graph to a pickle file for fast serialization.

    Args:
        output_path: Path to save the pickle file
    """
    # Prepare data for pickling
    data = {
        "schema_version": _SCHEMA_VERSION,
        "project_root": str(self.project_root) if self.project_root else None,
        "graph": self.graph,  # igraph objects are picklable
        "symbol_ranges": self.symbol_ranges,
        "name_to_vertex": self.name_to_vertex,
        "file_nodes": self._file_nodes,
        "file_edge_anchors": self._file_edge_anchors,
        "unified_to_names": self._unified_to_names,
    }

    with open(output_path, "wb") as f:
        pickle.dump(data, f, protocol=pickle.HIGHEST_PROTOCOL)

load_graph classmethod

load_graph(input_path)

Load a graph from a pickle file.

Parameters:

Name Type Description Default
input_path

Path to the pickle file

required

Returns:

Name Type Description
CodeGraph

Loaded graph instance

Source code in codenib/graph/code_graph.py
@classmethod
def load_graph(cls, input_path):
    """
    Load a graph from a pickle file.

    Args:
        input_path: Path to the pickle file

    Returns:
        CodeGraph: Loaded graph instance
    """
    with open(input_path, "rb") as f:
        data = compat_pickle.load(f)

    on_disk = data.get("schema_version")
    if on_disk != _SCHEMA_VERSION:
        raise ValueError(
            f"graph.pkl at {input_path} has schema_version={on_disk!r}, "
            f"expected {_SCHEMA_VERSION}. Delete the cached pickle and "
            "regenerate via the indexing pipeline (e.g. integration-serial)."
        )

    # Create new CodeGraph instance
    graph_instance = cls(project_root=data.get("project_root"))

    # Restore the igraph object and internal state
    graph_instance.graph = data["graph"]
    graph_instance.symbol_ranges = data.get("symbol_ranges", {})
    graph_instance.name_to_vertex = data.get("name_to_vertex", {})
    graph_instance._file_nodes = data.get("file_nodes", {})
    graph_instance._file_edge_anchors = data.get("file_edge_anchors", {})
    graph_instance._unified_to_names = data.get("unified_to_names", {})

    return graph_instance

get_graph

get_graph()

Get the igraph Graph object.

Returns:

Type Description

The igraph Graph instance

Source code in codenib/graph/code_graph.py
def get_graph(self):
    """
    Get the igraph Graph object.

    Returns:
        The igraph Graph instance
    """
    return self.graph

unified_index

unified_index()

Readable unified_name (full + bare) -> [canonical names], cached.

Rebuilds from vertex attributes when the persisted _unified_to_names is empty (older prebuilt graphs ship it empty even though every vertex carries a unified_name).

Source code in codenib/graph/code_graph.py
def unified_index(self):
    """Readable ``unified_name`` (full + bare) -> [canonical names], cached.

    Rebuilds from vertex attributes when the persisted ``_unified_to_names``
    is empty (older prebuilt graphs ship it empty even though every vertex
    carries a ``unified_name``).
    """
    cache = getattr(self, "_unified_index_cache", None)
    if cache is not None:
        return cache
    index = {}

    def _add(key, name):
        index.setdefault(key, []).append(name)

    pre = self._unified_to_names or {}
    if pre:
        for disp, names in pre.items():
            for nm in names:
                _add(disp, nm)
                _add(self._bare_symbol(disp), nm)
    else:
        for nm in self.name_to_vertex:
            info = self.get_node_info_by_name(nm) or {}
            u = info.get("unified_name")
            if u:
                _add(u, nm)
                _add(self._bare_symbol(u), nm)
    self._unified_index_cache = index
    return index

display_name

display_name(name)

Readable label for a canonical name (unified_name or itself).

Source code in codenib/graph/code_graph.py
def display_name(self, name):
    """Readable label for a canonical *name* (``unified_name`` or itself)."""
    info = self.get_node_info_by_name(name) or {}
    return info.get("unified_name") or name

resolve_symbol

resolve_symbol(symbol)

Map a readable/bare symbol to (canonical_name | None, candidates).

Tries the canonical name first, then the readable unified_name (full display, then bare token), then a substring match. None when ambiguous or missing; candidates are readable display names.

Source code in codenib/graph/code_graph.py
def resolve_symbol(self, symbol):
    """Map a readable/bare *symbol* to (canonical_name | None, candidates).

    Tries the canonical name first, then the readable ``unified_name``
    (full display, then bare token), then a substring match. ``None`` when
    ambiguous or missing; candidates are readable display names.
    """
    n2v = self.name_to_vertex or {}
    s = (symbol or "").strip().strip("`'\"")
    if not s:
        return None, []
    if s in n2v:
        return s, [self.display_name(s)]

    uni = self.unified_index()
    cands = []
    if uni:
        sbase = self._bare_symbol(s)
        for key in (s, s + "()", sbase):
            if key in uni:
                cands = list(dict.fromkeys(uni[key]))
                break
        if not cands:
            sub = []
            for disp, names in uni.items():
                if sbase and sbase.lower() in disp.lower():
                    sub += names
            cands = list(dict.fromkeys(sub))
    if not cands:  # languages whose canonical name is already readable
        base = s.split(".")[-1].split(":")[-1]
        cands = [k for k in n2v if k.endswith("." + s) or k.split(".")[-1] == base]
    canonical = cands[0] if len(cands) == 1 else None
    return canonical, [self.display_name(c) for c in cands[:8]]

get_node_info_by_name

get_node_info_by_name(node_name)

Get information about a node in the graph.

Parameters:

Name Type Description Default
node_name

Name of the node

required

Returns:

Type Description

Dictionary with vertex attributes or None if not found

Source code in codenib/graph/code_graph.py
def get_node_info_by_name(self, node_name):
    """
    Get information about a node in the graph.

    Args:
        node_name: Name of the node

    Returns:
        Dictionary with vertex attributes or None if not found
    """
    if isinstance(node_name, str):
        vertex = self.name_to_vertex.get(node_name)
        if vertex is None:
            return None

    return self.graph.vs[vertex].attributes()

get_node_info_by_id

get_node_info_by_id(node_id)

Get information about a node in the graph.

Parameters:

Name Type Description Default
node_id

ID of the node

required

Returns:

Type Description

Dictionary with vertex attributes or None if not found

Source code in codenib/graph/code_graph.py
def get_node_info_by_id(self, node_id):
    """
    Get information about a node in the graph.

    Args:
        node_id: ID of the node

    Returns:
        Dictionary with vertex attributes or None if not found
    """
    if isinstance(node_id, int):
        vertex = self.graph.vs[node_id]
        return vertex.attributes()

    return None

get_node_content

get_node_content(node_id)

Get the content of a node in the graph. Read the file from start_line to end_line. If the node is a file, return the file content.

Parameters:

Name Type Description Default
node_id

ID of the node

required

Returns:

Type Description

Content of the node or None if not found

Source code in codenib/graph/code_graph.py
def get_node_content(self, node_id):
    """
    Get the content of a node in the graph.
    Read the file from start_line to end_line.
    If the node is a file, return the file content.

    Args:
        node_id: ID of the node

    Returns:
        Content of the node or None if not found
    """
    node = self.graph.vs[node_id]
    node_type = node["type"] if "type" in node.attributes() else "unknown"
    if node_type == NODE_TYPE_FILE:
        # file path need to add self.project_root
        file_path = node["name"]
        if self.project_root:
            file_path = f"{self.project_root}/{file_path}"
        try:
            with open(file_path, "r") as f:
                content = f.read()
            return content
        except FileNotFoundError:
            print(f"File not found: {file_path}")
            return None
    elif is_symbol_node(node_type):
        # Graph source ranges are 0-based and inclusive.
        start_line = node["start_line"]
        end_line = node["end_line"]
        if (
            not isinstance(start_line, int)
            or isinstance(start_line, bool)
            or not isinstance(end_line, int)
            or isinstance(end_line, bool)
            or start_line < 0
            or end_line < start_line
        ):
            return None

        # Get the file path
        file_path = node["file"] if "file" in node.attributes() else None
        if self.project_root and file_path:
            file_path = f"{self.project_root}/{file_path}"

        # Read the file content
        if file_path:
            try:
                with open(file_path, "r") as f:
                    lines = f.readlines()
                return "".join(lines[start_line : end_line + 1])
            except FileNotFoundError:
                print(f"File not found: {file_path}")
                return None
    return None

get_neighbors

get_neighbors(node_name)

Get the neighbors of a node in the graph.

Parameters:

Name Type Description Default
node_name

Name of the node

required

Returns:

Type Description

List of neighbor vertex IDs

Source code in codenib/graph/code_graph.py
def get_neighbors(self, node_name):
    """
    Get the neighbors of a node in the graph.

    Args:
        node_name: Name of the node

    Returns:
        List of neighbor vertex IDs
    """
    if isinstance(node_name, str):
        vertex = self.name_to_vertex.get(node_name)
        if vertex is None:
            return []

    return self.graph.neighbors(vertex)

get_successors

get_successors(node_name, edge_types=None)

Get the successors of a node_name (outgoing edges).

Parameters:

Name Type Description Default
node_name

Name of the node

required
edge_types

Optional edge types to retain. Typed traversal returns unique neighbors even when the graph contains multi-edges.

None

Returns:

Type Description

List of successor vertex IDs

Source code in codenib/graph/code_graph.py
def get_successors(self, node_name, edge_types=None):
    """
    Get the successors of a node_name (outgoing edges).

    Args:
        node_name: Name of the node
        edge_types: Optional edge types to retain. Typed traversal returns
            unique neighbors even when the graph contains multi-edges.

    Returns:
        List of successor vertex IDs
    """
    if isinstance(node_name, str):
        vertex = self.name_to_vertex.get(node_name)
        if vertex is None:
            return []

    if edge_types is None:
        return self.graph.successors(vertex)
    return self._get_typed_adjacent(vertex, edge_types, mode="out")

get_predecessors

get_predecessors(node_name, edge_types=None)

Get the predecessors of a node_name (incoming edges).

Parameters:

Name Type Description Default
node_name

Name of the node

required
edge_types

Optional edge types to retain. Typed traversal returns unique neighbors even when the graph contains multi-edges.

None

Returns:

Type Description

List of predecessor vertex IDs

Source code in codenib/graph/code_graph.py
def get_predecessors(self, node_name, edge_types=None):
    """
    Get the predecessors of a node_name (incoming edges).

    Args:
        node_name: Name of the node
        edge_types: Optional edge types to retain. Typed traversal returns
            unique neighbors even when the graph contains multi-edges.

    Returns:
        List of predecessor vertex IDs
    """
    if isinstance(node_name, str):
        vertex = self.name_to_vertex.get(node_name)
        if vertex is None:
            return []

    if edge_types is None:
        return self.graph.predecessors(vertex)
    return self._get_typed_adjacent(vertex, edge_types, mode="in")

print_graph_basic_info

print_graph_basic_info()

Print basic information about the graph.

Source code in codenib/graph/code_graph.py
def print_graph_basic_info(self):
    """
    Print basic information about the graph.
    """
    print("Graph Summary:")
    print(f"  Number of vertices: {self.graph.vcount()}")
    print(f"  Number of edges: {self.graph.ecount()}")
    print(f"  Directed: {self.graph.is_directed()}")

    # Print vertex types
    vertex_types = self.graph.vs["type"]
    unique_vertex_types = set(vertex_types)
    print(f"  Unique vertex types: {unique_vertex_types}")

    # Print edge types
    edge_types = self.graph.es["type"]
    unique_edge_types = set(edge_types)
    print(f"  Unique edge types: {unique_edge_types}")

visualize_graph

visualize_graph(output_path=None, width=800, height=600, layout='fruchterman_reingold')

Visualize the code graph with different colors for different node and edge types.

Parameters:

Name Type Description Default
output_path

Path to save the visualization (optional)

None
width

Width of the plot in pixels

800
height

Height of the plot in pixels

600
layout

Layout algorithm to use ('fruchterman_reingold', 'kk', 'grid', etc.)

'fruchterman_reingold'

Returns:

Type Description

A matplotlib figure object if output_path is not provided

Source code in codenib/graph/code_graph.py
def visualize_graph(
    self, output_path=None, width=800, height=600, layout="fruchterman_reingold"
):
    """
    Visualize the code graph with different colors for different node and edge types.

    Args:
        output_path: Path to save the visualization (optional)
        width: Width of the plot in pixels
        height: Height of the plot in pixels
        layout: Layout algorithm to use ('fruchterman_reingold', 'kk', 'grid', etc.)

    Returns:
        A matplotlib figure object if output_path is not provided
    """
    if self.graph.vcount() == 0:
        print("Graph is empty. Nothing to visualize.")
        return None

    try:
        import matplotlib.pyplot as plt
    except ImportError as exc:
        raise RuntimeError(
            "static graph rendering requires matplotlib; install it separately"
        ) from exc

    # Define color schemes
    node_type_colors = {
        NODE_TYPE_FILE: "skyblue",
        NODE_TYPE_SYMBOL: "lightgreen",
        NODE_TYPE_CLASS: "gold",
        NODE_TYPE_FUNCTION: "lightgreen",
        NODE_TYPE_METHOD: "lightcoral",
        NODE_TYPE_FIELD: "plum",
        NODE_TYPE_DIRECTORY: "orange",  # Add directory node color
        "root": "lightgrey",  # Root node color
        # Add more node types and colors as needed
    }

    edge_type_colors = {
        EDGE_TYPE_REFERENCE: "red",
        EDGE_TYPE_CONTAIN: "blue",
        # Add more edge types and colors as needed
    }

    # Define visual style
    visual_style = {}

    # Set vertex colors based on type
    vertex_colors = []
    for vertex in self.graph.vs:
        node_type = vertex["type"] if "type" in vertex.attributes() else "unknown"
        vertex_colors.append(node_type_colors.get(node_type, "grey"))
    visual_style["vertex_color"] = vertex_colors

    # Set vertex labels
    visual_style["vertex_label"] = [
        v["name"].split("/")[-1] if "/" in v["name"] else v["name"]
        for v in self.graph.vs
    ]
    visual_style["vertex_label_size"] = 8

    # Set vertex sizes (files can be larger than symbols)
    vertex_sizes = []
    for vertex in self.graph.vs:
        if "type" in vertex.attributes() and vertex["type"] == NODE_TYPE_FILE:
            vertex_sizes.append(20)
        else:
            vertex_sizes.append(10)
    visual_style["vertex_size"] = vertex_sizes

    # Set edge colors based on type
    edge_colors = []
    for edge in self.graph.es:
        edge_type = edge["type"] if "type" in edge.attributes() else "unknown"
        edge_colors.append(edge_type_colors.get(edge_type, "grey"))
    visual_style["edge_color"] = edge_colors

    # Set edge width
    visual_style["edge_width"] = 1.0

    # Calculate layout
    if layout == "fruchterman_reingold":
        visual_style["layout"] = self.graph.layout_fruchterman_reingold()
    elif layout == "kk":
        visual_style["layout"] = self.graph.layout_kamada_kawai()
    elif layout == "grid":
        visual_style["layout"] = self.graph.layout_grid()
    else:
        visual_style["layout"] = self.graph.layout(layout)

    # Adjust dimensions
    visual_style["bbox"] = (width, height)
    visual_style["margin"] = 40

    # Create legend data
    legend_data = {
        "Node Types": [
            (color, node_type) for node_type, color in node_type_colors.items()
        ],
        "Edge Types": [
            (color, edge_type) for edge_type, color in edge_type_colors.items()
        ],
    }

    # Create the plot
    fig, ax = plt.subplots(figsize=(width / 100, height / 100))

    # Plot the graph
    ig.plot(self.graph, target=ax, **visual_style)

    # Add legend for node types
    node_legend_patches = []
    for color, label in legend_data["Node Types"]:
        node_legend_patches.append(
            plt.Line2D(
                [0],
                [0],
                marker="o",
                color="w",
                markerfacecolor=color,
                markersize=10,
                label=label,
            )
        )

    # Add legend for edge types
    edge_legend_patches = []
    for color, label in legend_data["Edge Types"]:
        edge_legend_patches.append(
            plt.Line2D([0], [0], color=color, lw=2, label=label)
        )

    # Add legends to plot
    ax.legend(
        handles=node_legend_patches + edge_legend_patches,
        loc="upper right",
        title="Legend",
        frameon=True,
    )

    # Save or show the plot
    if output_path:
        plt.savefig(output_path, dpi=100, bbox_inches="tight")
        print(f"Graph visualization saved to {output_path}")
        plt.close(fig)
        return None
    else:
        plt.tight_layout()
        return fig

GraphPatcher

GraphPatcher(
    project_root: str, code_graph: CodeGraph | None, language: str, mode: str = "symbol"
)

Orchestrates incremental graph updates using LSP.

Thin router that delegates to language-specific PatcherBase subclasses.

Usage::

graph = CodeGraph.load_graph("graph.pkl")
patcher = GraphPatcher("/path/to/project", graph, "rust")
changed = GraphPatcher.detect_changed_files(
    "/path/to/project", "abc123", "HEAD"
)
stats = patcher.patch_files(changed)
graph.save_graph("graph.pkl")

Methods:

Name Description
start_lsp

Start the LSP server for warm-up.

stop_lsp

Stop the LSP server.

patch_files

Apply incremental patch for all changed files.

detect_changed_files

Detect changed files between two commits.

get_changed_line_ranges

Get line ranges changed between two commits.

Source code in codenib/graph/incremental/graph_patcher.py
def __init__(
    self,
    project_root: str,
    code_graph: Optional[CodeGraph],
    language: str,
    mode: str = "symbol",
):
    self.project_root = project_root
    self.language = language
    self.mode = mode
    self._impl: PatcherBase = _create_patcher(
        language, project_root, code_graph, mode=mode
    )

start_lsp

start_lsp(skip_probe: bool = False)

Start the LSP server for warm-up.

Source code in codenib/graph/incremental/graph_patcher.py
def start_lsp(self, skip_probe: bool = False):
    """Start the LSP server for warm-up."""
    self._impl.start_lsp(skip_probe=skip_probe)

stop_lsp

stop_lsp()

Stop the LSP server.

Source code in codenib/graph/incremental/graph_patcher.py
def stop_lsp(self):
    """Stop the LSP server."""
    self._impl.stop_lsp()

patch_files

patch_files(changed_files: dict, **kwargs) -> dict

Apply incremental patch for all changed files.

Source code in codenib/graph/incremental/graph_patcher.py
def patch_files(self, changed_files: dict, **kwargs) -> dict:
    """Apply incremental patch for all changed files."""
    return self._impl.patch_files(changed_files, **kwargs)

detect_changed_files staticmethod

detect_changed_files(
    project_root: str,
    base_commit: str,
    target_commit: str = "HEAD",
    extensions: set[str] | None = None,
) -> dict

Detect changed files between two commits.

Source code in codenib/graph/incremental/graph_patcher.py
@staticmethod
def detect_changed_files(
    project_root: str,
    base_commit: str,
    target_commit: str = "HEAD",
    extensions: Optional[set[str]] = None,
) -> dict:
    """Detect changed files between two commits."""
    return change_mgr.detect_changed_files(
        project_root, base_commit, target_commit, extensions
    )

get_changed_line_ranges staticmethod

get_changed_line_ranges(
    project_root: str, file_path: str, base_commit: str, target_commit: str = "HEAD"
) -> list[tuple[int, int, int, int]]

Get line ranges changed between two commits.

Each hunk is reported as (old_start, old_end, new_start, new_end) in 0-indexed inclusive form so callers can compare line counts.

Source code in codenib/graph/incremental/graph_patcher.py
@staticmethod
def get_changed_line_ranges(
    project_root: str,
    file_path: str,
    base_commit: str,
    target_commit: str = "HEAD",
) -> list[tuple[int, int, int, int]]:
    """Get line ranges changed between two commits.

    Each hunk is reported as (old_start, old_end, new_start, new_end)
    in 0-indexed inclusive form so callers can compare line counts.
    """
    return change_mgr.get_changed_line_ranges(
        project_root, file_path, base_commit, target_commit
    )

LSPClient

LSPClient(
    command: list[str],
    project_root: str,
    language: str,
    init_options: dict | None = None,
)

Synchronous LSP client that communicates with a language server via stdio.

Usage::

with LSPClient(["rust-analyzer"], "/path/to/project", "rust") as client:
    symbols = client.document_symbol("src/main.rs")
    refs = client.references("src/main.rs", 10, 4)

Methods:

Name Description
start

Start the LSP server and send initialize/initialized.

shutdown

Gracefully shut down the LSP server.

open_document

Send textDocument/didOpen for a file (idempotent).

close_document

Send textDocument/didClose for a file.

wait_for_analysis

Wait for the LSP server to finish analyzing a file.

document_symbol

Get hierarchical document symbols for a file.

semantic_tokens_full

Get semantic tokens for the entire file.

semantic_tokens_range

Get semantic tokens for a specific line range.

decode_semantic_tokens

Decode semantic tokens response into a list of token dicts.

drain_notifications

Drain any pending notifications without waiting for a response.

wait_until_idle

Wait until no $/progress tokens are active for idle_grace_s.

get_lsp_command

Get the default LSP server command for a language.

check_lsp_available

Check if the LSP server binary is available.

Source code in codenib/graph/incremental/lsp_client.py
def __init__(
    self,
    command: list[str],
    project_root: str,
    language: str,
    init_options: Optional[dict] = None,
):
    self.command = command
    self.project_root = Path(project_root).resolve()
    self.root_uri = self.project_root.as_uri()
    self.language = language
    self.init_options = init_options
    self.process: Optional[subprocess.Popen] = None
    self._next_id = 1
    self._opened_files: set[str] = set()

    # Populated during start() from server capabilities
    self.semantic_tokens_legend: Optional[dict] = None
    self.supports_semantic_tokens_range: bool = False

    # Active $/progress tokens (set by _handle_progress, cleared by
    # 'end' notifications). wait_until_idle() polls this for emptiness
    # to know when LSP background work has finished.
    self._active_progress: dict = {}

    # Pending request tracking: limit in-flight requests to avoid
    # overwhelming the LSP server (clangd hangs after ~570 queued requests)
    self._pending_ids: set[int] = set()
    self._max_pending: int = 10

    # Buffer for out-of-order responses: when _request reads a response
    # for a different request ID, it stores it here instead of discarding.
    self._response_buffer: dict[int, dict] = {}

    # Track the last error from _request so callers can decide whether
    # to retry (transient errors like -32801) or not (permanent errors).
    self._last_error: Optional[dict] = None

start

start(skip_probe: bool = False) -> dict

Start the LSP server and send initialize/initialized.

Parameters:

Name Type Description Default
skip_probe bool

If True, skip the blocking readiness probe. Use this when starting early for warm-up — the server will analyze the project in the background.

False
Source code in codenib/graph/incremental/lsp_client.py
def start(self, skip_probe: bool = False) -> dict:
    """Start the LSP server and send initialize/initialized.

    Args:
        skip_probe: If True, skip the blocking readiness probe.
            Use this when starting early for warm-up — the server
            will analyze the project in the background.
    """
    logger.info(f"Starting LSP server: {' '.join(self.command)}")
    env = _lsp_process_env(self.language, self.project_root)
    self.process = subprocess.Popen(
        self.command,
        stdin=subprocess.PIPE,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        cwd=str(self.project_root),
        env=env,
    )

    # Build capabilities we request
    capabilities = {
        "textDocument": {
            "documentSymbol": {
                "hierarchicalDocumentSymbolSupport": True,
            },
            "references": {},
            "definition": {
                "linkSupport": True,
            },
            "semanticTokens": {
                "requests": {"full": True},
                "tokenTypes": [],
                "tokenModifiers": [],
                "formats": ["relative"],
            },
        },
        "window": {
            "workDoneProgress": True,
        },
    }

    # Language-specific initialization options
    init_opts = dict(self.init_options) if self.init_options else {}
    if self.language == "go":
        # gopls disables semanticTokens by default — enable it
        init_opts.setdefault("semanticTokens", True)

    params = {
        "processId": os.getpid(),
        "rootUri": self.root_uri,
        "rootPath": str(self.project_root),
        "capabilities": capabilities,
        "workspaceFolders": [
            {"uri": self.root_uri, "name": self.project_root.name}
        ],
    }
    if init_opts:
        params["initializationOptions"] = init_opts

    result = self._request("initialize", params, timeout=120)
    self._notify("initialized", {})

    # Extract semantic tokens legend from server capabilities
    if result:
        sem_provider = result.get("capabilities", {}).get(
            "semanticTokensProvider", {}
        )
        if isinstance(sem_provider, dict):
            self.semantic_tokens_legend = sem_provider.get("legend")
            self.supports_semantic_tokens_range = bool(sem_provider.get("range"))

    # Wait for server readiness (unless skipped for warm-up)
    if not skip_probe:
        self._wait_for_ready()

    logger.info("LSP server initialized successfully")
    return result or {}

shutdown

shutdown()

Gracefully shut down the LSP server.

Source code in codenib/graph/incremental/lsp_client.py
def shutdown(self):
    """Gracefully shut down the LSP server."""
    if self.process is None:
        return
    try:
        self._request("shutdown", None, timeout=10)
        self._notify("exit", None)
        self.process.wait(timeout=5)
    except Exception:
        if self.process:
            self.process.kill()
            self.process.wait()
    finally:
        self.process = None
        self._opened_files.clear()

open_document

open_document(file_path: str)

Send textDocument/didOpen for a file (idempotent).

Source code in codenib/graph/incremental/lsp_client.py
def open_document(self, file_path: str):
    """Send textDocument/didOpen for a file (idempotent)."""
    abs_path = self._abs_path(file_path)
    uri = Path(abs_path).as_uri()
    if uri in self._opened_files:
        return
    try:
        text = Path(abs_path).read_text(errors="replace")
    except FileNotFoundError:
        logger.warning(f"Cannot open {abs_path}: file not found")
        return

    self._notify(
        "textDocument/didOpen",
        {
            "textDocument": {
                "uri": uri,
                "languageId": self._detect_language_id(abs_path),
                "version": 1,
                "text": text,
            },
        },
    )
    self._opened_files.add(uri)

close_document

close_document(file_path: str)

Send textDocument/didClose for a file.

Source code in codenib/graph/incremental/lsp_client.py
def close_document(self, file_path: str):
    """Send textDocument/didClose for a file."""
    abs_path = self._abs_path(file_path)
    uri = Path(abs_path).as_uri()
    if uri not in self._opened_files:
        return
    self._notify(
        "textDocument/didClose",
        {
            "textDocument": {"uri": uri},
        },
    )
    self._opened_files.discard(uri)

wait_for_analysis

wait_for_analysis(file_path: str, max_wait: float = 15.0, poll_interval: float = 0.5)

Wait for the LSP server to finish analyzing a file.

After didOpen/didChange, servers like pyright need time to re-analyze before references() returns correct results. We poll with a lightweight hover request until it responds quickly, indicating analysis is complete.

Source code in codenib/graph/incremental/lsp_client.py
def wait_for_analysis(
    self,
    file_path: str,
    max_wait: float = 15.0,
    poll_interval: float = 0.5,
):
    """Wait for the LSP server to finish analyzing a file.

    After didOpen/didChange, servers like pyright need time to
    re-analyze before references() returns correct results.
    We poll with a lightweight hover request until it responds
    quickly, indicating analysis is complete.
    """
    abs_path = self._abs_path(file_path)
    uri = Path(abs_path).as_uri()
    start = time.time()
    while time.time() - start < max_wait:
        t0 = time.time()
        try:
            self._request(
                "textDocument/hover",
                {
                    "textDocument": {"uri": uri},
                    "position": {"line": 0, "character": 0},
                },
                timeout=5,
            )
        except Exception as exc:
            logger.debug(f"Analysis probe failed: {exc}")
        elapsed = time.time() - t0
        # If hover responds in < 1s, server is likely done analyzing
        if elapsed < 1.0:
            return
        time.sleep(poll_interval)

document_symbol

document_symbol(
    file_path: str, retries: int = 0, retry_delay: float = 0.5
) -> list[dict]

Get hierarchical document symbols for a file.

Default retries=0 — null is a legitimate "no symbols" answer per LSP spec. Previous default of 5 retries × 5s sleep wasted up to 25s/call when servers returned null. If the server is still loading at warmup time, callers should wait at the call-site, not inside this primitive.

Source code in codenib/graph/incremental/lsp_client.py
def document_symbol(
    self, file_path: str, retries: int = 0, retry_delay: float = 0.5
) -> list[dict]:
    """Get hierarchical document symbols for a file.

    Default ``retries=0`` — null is a legitimate "no symbols" answer
    per LSP spec. Previous default of 5 retries × 5s sleep wasted up
    to 25s/call when servers returned null. If the server is still
    loading at warmup time, callers should wait at the call-site,
    not inside this primitive.
    """
    t0 = time.monotonic()
    self.open_document(file_path)
    abs_path = self._abs_path(file_path)
    uri = Path(abs_path).as_uri()

    for attempt in range(retries + 1):
        result = self._request(
            "textDocument/documentSymbol",
            {
                "textDocument": {"uri": uri},
            },
            timeout=10,  # longest non-empty observed: 138ms
        )
        if result is not None:
            _profile_log(
                "documentSymbol",
                file_path,
                None,
                None,
                time.monotonic() - t0,
                len(result) if isinstance(result, list) else 0,
            )
            return result
        if attempt < retries:
            logger.debug(
                f"documentSymbol returned null for {file_path}, "
                f"retrying in {retry_delay}s ({attempt + 1}/{retries})"
            )
            time.sleep(retry_delay)

    _profile_log("documentSymbol", file_path, None, None, time.monotonic() - t0, 0)
    return []

semantic_tokens_full

semantic_tokens_full(file_path: str, timeout: float = 15) -> dict | None

Get semantic tokens for the entire file.

Returns raw response with data field (delta-encoded integers). Use decode_semantic_tokens() to decode. Returns None on timeout or error (details logged by _request).

Default timeout 15s: a server that needs longer is pathological (basedpyright on sklearn's test_pls.py used to take 60s). The patcher tolerates None — that file simply skips outgoing-ref discovery for this run. 60s × N pathological files dominated runtime before this cap.

Source code in codenib/graph/incremental/lsp_client.py
def semantic_tokens_full(
    self, file_path: str, timeout: float = 15
) -> Optional[dict]:
    """Get semantic tokens for the entire file.

    Returns raw response with ``data`` field (delta-encoded integers).
    Use ``decode_semantic_tokens()`` to decode.
    Returns None on timeout or error (details logged by _request).

    Default timeout 15s: a server that needs longer is pathological
    (basedpyright on sklearn's test_pls.py used to take 60s). The
    patcher tolerates None — that file simply skips outgoing-ref
    discovery for this run. 60s × N pathological files dominated
    runtime before this cap.
    """
    t0 = time.monotonic()
    self.open_document(file_path)
    abs_path = self._abs_path(file_path)
    uri = Path(abs_path).as_uri()
    result = self._request(
        "textDocument/semanticTokens/full",
        {
            "textDocument": {"uri": uri},
        },
        timeout=timeout,
    )
    n_tokens = 0
    if result is not None:
        n_tokens = len(result.get("data", [])) // 5
        logger.debug(f"semanticTokens for {file_path}: {n_tokens} tokens")
    _profile_log(
        "semanticTokens/full",
        file_path,
        None,
        None,
        time.monotonic() - t0,
        n_tokens,
    )
    return result

semantic_tokens_range

semantic_tokens_range(
    file_path: str, start_line: int, end_line: int, timeout: float = 30
) -> dict | None

Get semantic tokens for a specific line range.

Faster than full for large files when only a small range is needed. Returns None if not supported or on error.

Source code in codenib/graph/incremental/lsp_client.py
def semantic_tokens_range(
    self,
    file_path: str,
    start_line: int,
    end_line: int,
    timeout: float = 30,
) -> Optional[dict]:
    """Get semantic tokens for a specific line range.

    Faster than full for large files when only a small range is needed.
    Returns None if not supported or on error.
    """
    t0 = time.monotonic()
    self.open_document(file_path)
    abs_path = self._abs_path(file_path)
    uri = Path(abs_path).as_uri()
    result = self._request(
        "textDocument/semanticTokens/range",
        {
            "textDocument": {"uri": uri},
            "range": {
                "start": {"line": start_line, "character": 0},
                "end": {"line": end_line + 1, "character": 0},
            },
        },
        timeout=timeout,
    )
    n_tokens = 0
    if result is not None:
        n_tokens = len(result.get("data", [])) // 5
        logger.debug(
            f"semanticTokens/range for {file_path} "
            f"L{start_line}-{end_line}: {n_tokens} tokens"
        )
    _profile_log(
        "semanticTokens/range",
        file_path,
        start_line,
        end_line,
        time.monotonic() - t0,
        n_tokens,
    )
    return result

decode_semantic_tokens

decode_semantic_tokens(tokens_response: dict, file_path: str) -> list[dict]

Decode semantic tokens response into a list of token dicts.

Each token dict has: line, character, length, token_type, modifiers, text.

Source code in codenib/graph/incremental/lsp_client.py
def decode_semantic_tokens(
    self, tokens_response: dict, file_path: str
) -> list[dict]:
    """Decode semantic tokens response into a list of token dicts.

    Each token dict has: line, character, length, token_type, modifiers, text.
    """
    if not tokens_response or "data" not in tokens_response:
        return []
    if not self.semantic_tokens_legend:
        logger.warning("No semantic tokens legend available")
        return []

    data = tokens_response["data"]
    token_types = self.semantic_tokens_legend.get("tokenTypes", [])
    token_modifiers = self.semantic_tokens_legend.get("tokenModifiers", [])

    # Read file content for extracting token text
    abs_path = self._abs_path(file_path)
    try:
        lines = Path(abs_path).read_text(errors="replace").splitlines()
    except FileNotFoundError:
        lines = []

    tokens = []
    current_line = 0
    current_char = 0

    for i in range(0, len(data), 5):
        if i + 4 >= len(data):
            break
        delta_line = data[i]
        delta_char = data[i + 1]
        length = data[i + 2]
        type_idx = data[i + 3]
        mod_bits = data[i + 4]

        if delta_line > 0:
            current_line += delta_line
            current_char = delta_char
        else:
            current_char += delta_char

        # Decode type
        type_name = (
            token_types[type_idx]
            if type_idx < len(token_types)
            else f"type_{type_idx}"
        )

        # Decode modifiers (bitmask)
        mods = []
        for bit_idx, mod_name in enumerate(token_modifiers):
            if mod_bits & (1 << bit_idx):
                mods.append(mod_name)

        # Extract text from source
        text = ""
        if current_line < len(lines):
            line_text = lines[current_line]
            text = line_text[current_char : current_char + length]

        tokens.append(
            {
                "line": current_line,
                "character": current_char,
                "length": length,
                "token_type": type_name,
                "modifiers": mods,
                "text": text,
            }
        )

    return tokens

drain_notifications

drain_notifications(timeout_s: float = 0.5) -> int

Drain any pending notifications without waiting for a response.

Returns the number of messages drained. Used by wait_until_idle to keep the progress map fresh while polling.

Source code in codenib/graph/incremental/lsp_client.py
def drain_notifications(self, timeout_s: float = 0.5) -> int:
    """Drain any pending notifications without waiting for a response.

    Returns the number of messages drained. Used by ``wait_until_idle``
    to keep the progress map fresh while polling.
    """
    n = 0
    deadline = time.monotonic() + timeout_s
    while time.monotonic() < deadline:
        remaining = deadline - time.monotonic()
        if remaining <= 0:
            break
        ready, _, _ = select.select(
            [self.process.stdout], [], [], min(remaining, 0.05)
        )
        if not ready:
            break
        msg = self._read_message(timeout=remaining)
        if msg is None:
            break
        n += 1
        # We don't dispatch unmatched responses here; just count and
        # let _handle_progress (called inside _read_message) update state.
    return n

wait_until_idle

wait_until_idle(max_wait_s: float = 60.0, idle_grace_s: float = 1.0) -> bool

Wait until no $/progress tokens are active for idle_grace_s.

Returns True if the server became idle within max_wait_s, False on timeout. Useful before timed regions so background indexing (rust-analyzer cache priming, gopls workspace load, basedpyright type analysis) doesn't bleed into measurements.

Source code in codenib/graph/incremental/lsp_client.py
def wait_until_idle(
    self, max_wait_s: float = 60.0, idle_grace_s: float = 1.0
) -> bool:
    """Wait until no ``$/progress`` tokens are active for ``idle_grace_s``.

    Returns True if the server became idle within ``max_wait_s``,
    False on timeout. Useful before timed regions so background
    indexing (rust-analyzer cache priming, gopls workspace load,
    basedpyright type analysis) doesn't bleed into measurements.
    """
    deadline = time.monotonic() + max_wait_s
    idle_since: Optional[float] = None
    while time.monotonic() < deadline:
        self.drain_notifications(timeout_s=0.2)
        if not self._active_progress:
            if idle_since is None:
                idle_since = time.monotonic()
            elif time.monotonic() - idle_since >= idle_grace_s:
                return True
        else:
            idle_since = None
        time.sleep(0.1)
    return False

get_lsp_command staticmethod

get_lsp_command(language: str) -> list[str] | None

Get the default LSP server command for a language.

Source code in codenib/graph/incremental/lsp_client.py
@staticmethod
def get_lsp_command(language: str) -> Optional[list[str]]:
    """Get the default LSP server command for a language."""
    cmd = lsp_command_for_language(language)
    if not cmd:
        return None
    binary = cmd[0]
    resolved = resolve_lsp_binary(binary)
    if resolved:
        return [resolved] + cmd[1:]
    return cmd

check_lsp_available staticmethod

check_lsp_available(language: str) -> bool

Check if the LSP server binary is available.

Source code in codenib/graph/incremental/lsp_client.py
@staticmethod
def check_lsp_available(language: str) -> bool:
    """Check if the LSP server binary is available."""
    cmd = LSPClient.get_lsp_command(language)
    if not cmd:
        return False
    if normalize_graph_language(language) == "rust":
        try:
            subprocess.run(
                cmd + ["--version"],
                check=True,
                capture_output=True,
                text=True,
                timeout=10,
            )
            return True
        except (
            subprocess.CalledProcessError,
            FileNotFoundError,
            subprocess.TimeoutExpired,
        ):
            return False
    return Path(cmd[0]).exists() or resolve_lsp_binary(cmd[0]) is not None

PatcherBase

PatcherBase(
    project_root: str,
    code_graph: CodeGraph,
    lsp_client: LSPClient | None = None,
    mode: str = "symbol",
)

Bases: SubgraphMgr

Base class for language-specific incremental patchers.

Inherits SubgraphMgr for subgraph operations. Adds: - patch_files: top-level entry point - _rebuild_prepare_vertices / _rebuild_connect_edges: full rebuild for added/renamed - _incremental_prepare_vertices / _incremental_connect_edges: symbol-level diff - _classify_symbols, _remap_severed_edges, etc.

Subclasses must implement: - _build_unified_name: construct unified_name from LSP symbol info - _get_crossfile_token_types: which semanticToken types to query

Subclasses may override: - get_lsp_command: custom LSP server command - _language_id: custom LSP language id - flatten_symbols: convert LSP documentSymbol to flat dict (default provided) - get_old_symbols: extract old symbols from graph (default provided) - GRAPH_SYMBOL_KINDS: which SymbolKinds to process (default provided)

  • "symbol" (default): the actual symbol-level incremental patcher.
  • "naive": file-granularity baseline used by the sequential benchmark (issue #129). For modified files, every old symbol is classified as deleted and every new symbol as added — no shifted / affected / unchanged / invisible buckets. Untouched files are still left alone. Only the classifier behaviour changes; the rest of the pipeline (LSP queries, edge reconnection) is identical to symbol mode, which is exactly what makes it the right A/B comparison.

Methods:

Name Description
get_lsp_command

Return the LSP server command for this language.

flatten_symbols

Convert LSP documentSymbol to flat {unified_name: metadata} dict.

get_old_symbols

Extract existing definition vertices from the graph for a file.

start_lsp

Start the LSP server, resolving binary path.

stop_lsp

Stop the LSP server.

patch_files

Apply incremental patch for all changed files.

Source code in codenib/graph/incremental/patcher_base.py
def __init__(
    self,
    project_root: str,
    code_graph: CodeGraph,
    lsp_client: Optional[LSPClient] = None,
    mode: str = "symbol",
):
    """``mode`` is one of:

    - ``"symbol"`` (default): the actual symbol-level incremental patcher.
    - ``"naive"``: file-granularity baseline used by the sequential
      benchmark (issue #129). For modified files, every old symbol is
      classified as deleted and every new symbol as added — no shifted /
      affected / unchanged / invisible buckets. Untouched files are still
      left alone. Only the classifier behaviour changes; the rest of the
      pipeline (LSP queries, edge reconnection) is identical to symbol
      mode, which is exactly what makes it the right A/B comparison.
    """
    super().__init__(project_root, code_graph, lsp_client)
    if mode not in ("symbol", "naive"):
        raise ValueError(f"unknown patcher mode: {mode!r}")
    self.mode = mode
    self.profiler = Profiler()

get_lsp_command

get_lsp_command() -> list[str]

Return the LSP server command for this language.

Source code in codenib/graph/incremental/patcher_base.py
def get_lsp_command(self) -> list[str]:
    """Return the LSP server command for this language."""
    if not self.REGISTRY_LANGUAGE:
        raise NotImplementedError(
            f"{self.__class__.__name__} must define REGISTRY_LANGUAGE "
            "or override get_lsp_command()."
        )
    command = lsp_command_for_language(self.REGISTRY_LANGUAGE)
    if command is None:
        raise ValueError(
            f"No LSP command registered for {self.REGISTRY_LANGUAGE!r}"
        )
    return command

flatten_symbols abstractmethod

flatten_symbols(file_path: str, lsp_symbols: list[dict]) -> dict[str, dict]

Convert LSP documentSymbol to flat {unified_name: metadata} dict.

Must filter by GRAPH_SYMBOL_KINDS and handle language-specific naming (e.g. Rust impl blocks, Go pointer receivers).

Source code in codenib/graph/incremental/patcher_base.py
@abstractmethod
def flatten_symbols(
    self, file_path: str, lsp_symbols: list[dict]
) -> dict[str, dict]:
    """Convert LSP documentSymbol to flat {unified_name: metadata} dict.

    Must filter by GRAPH_SYMBOL_KINDS and handle language-specific
    naming (e.g. Rust impl blocks, Go pointer receivers).
    """

get_old_symbols

get_old_symbols(file_path: str) -> dict[str, dict]

Extract existing definition vertices from the graph for a file.

Default implementation: returns symbols with valid start_line (skips SCIP reference-only vertices). Override for language-specific filtering.

Source code in codenib/graph/incremental/patcher_base.py
def get_old_symbols(self, file_path: str) -> dict[str, dict]:
    """Extract existing definition vertices from the graph for a file.

    Default implementation: returns symbols with valid start_line
    (skips SCIP reference-only vertices).
    Override for language-specific filtering.
    """
    g = self.code_graph
    result = {}
    vids = getattr(g, "file_to_vertices", {}).get(file_path, set())
    for vid in vids:
        v = g.graph.vs[vid]
        attrs = v.attributes()
        uname = attrs.get("unified_name")
        if not uname:
            continue
        sl = attrs.get("start_line")
        if sl is None:
            continue  # SCIP reference vertex, not a definition
        result[uname] = {
            "vertex_name": v["name"],
            "start_line": sl,
            "end_line": attrs.get("end_line", sl),
        }
    return result

start_lsp

start_lsp(skip_probe: bool = False)

Start the LSP server, resolving binary path.

Source code in codenib/graph/incremental/patcher_base.py
def start_lsp(self, skip_probe: bool = False):
    """Start the LSP server, resolving binary path."""
    from .lsp_client import resolve_lsp_binary

    cmd = self.get_lsp_command()
    resolved = resolve_lsp_binary(cmd[0])
    if resolved:
        cmd = [resolved] + cmd[1:]

    self.lsp_client = LSPClient(cmd, str(self.project_root), self._language_id())
    self.lsp_client.start(skip_probe=skip_probe)

stop_lsp

stop_lsp()

Stop the LSP server.

Source code in codenib/graph/incremental/patcher_base.py
def stop_lsp(self):
    """Stop the LSP server."""
    if self.lsp_client:
        self.lsp_client.shutdown()
        self.lsp_client = None

patch_files

patch_files(
    changed_files: dict, earlier_commit: str = None, later_commit: str = None
) -> dict

Apply incremental patch for all changed files.

Parameters:

Name Type Description Default
changed_files dict

dict with modified/added/deleted/renamed lists.

required
earlier_commit str

base commit (required for modified files).

None
later_commit str

target commit.

None

Returns:

Type Description
dict

Stats dict.

Source code in codenib/graph/incremental/patcher_base.py
def patch_files(
    self,
    changed_files: dict,
    earlier_commit: str = None,
    later_commit: str = None,
) -> dict:
    """Apply incremental patch for all changed files.

    Args:
        changed_files: dict with modified/added/deleted/renamed lists.
        earlier_commit: base commit (required for modified files).
        later_commit: target commit.

    Returns:
        Stats dict.
    """
    effective_later_commit = later_commit or "HEAD"

    # Drop files the language excludes from indexing (e.g. _test.go for
    # Go, where SCIP-go skips test files by design — see
    # scip_decode_go.py).
    changed_files = self._filter_changed_files(changed_files)

    # Auto-start LSP if not already running
    if self.lsp_client is None:
        self.start_lsp()

    # Rebuild indexes (vertex IDs may have changed from previous patch)
    self.build_indexes()

    total_stats = {
        "files_deleted": 0,
        "files_modified": 0,
        "files_added": 0,
        "files_renamed": 0,
        "vertices_deleted": 0,
        "vertices_created": 0,
        "vertices_shifted": 0,
        "refs_incoming": 0,
        "refs_outgoing": 0,
        "refs_remapped": 0,
        "refs_unmatched": 0,
    }

    # Clear stale caches from previous patch_files calls
    self._semantic_tokens_cache.clear()

    nodes_before = self.code_graph.graph.vcount()
    edges_before = self.code_graph.graph.ecount()

    # ── Round 1: Build all vertices ──────────────────────────
    # Process deletions, then build vertices for added/renamed/modified.
    # No edge connections yet — ensures all vertices exist first.

    # 1a. Deletions
    for path in changed_files.get("deleted", []):
        with self.profiler.section("delete_subgraph"):
            self.delete_file_subgraph(path)
        total_stats["files_deleted"] += 1

    # 1b. Renames: delete old
    for old_path, _ in changed_files.get("renamed", []):
        with self.profiler.section("delete_subgraph"):
            self.delete_file_subgraph(old_path)

    # 1c. Added files: build vertices
    add_contexts = []
    for path in changed_files.get("added", []):
        ctx = self._rebuild_prepare_vertices(path, is_new=True)
        add_contexts.append((path, ctx))
        self._merge_stats(total_stats, ctx["file_stats"])
        total_stats["files_added"] += 1

    # 1d. Renamed files: build vertices for new path
    rename_contexts = []
    for _, new_path in changed_files.get("renamed", []):
        ctx = self._rebuild_prepare_vertices(new_path, is_new=True)
        rename_contexts.append((new_path, ctx))
        self._merge_stats(total_stats, ctx["file_stats"])
        total_stats["files_renamed"] += 1

    # 1e. Modified files: classify + build vertices
    modified = changed_files.get("modified", [])
    if modified and not earlier_commit:
        raise ValueError("earlier_commit required for modified files")
    mod_contexts = []
    for path in modified:
        ctx = self._incremental_prepare_vertices(
            path,
            earlier_commit,
            effective_later_commit,
        )
        if ctx is not None:
            mod_contexts.append((path, ctx))
            self._merge_stats(total_stats, ctx["file_stats"])
        total_stats["files_modified"] += 1

    # ── Round 2: Connect all edges ───────────────────────────
    # All vertices from all files now exist.

    for path, ctx in add_contexts:
        edge_stats = self._rebuild_connect_edges(path, ctx)
        self._merge_stats(total_stats, edge_stats)

    for path, ctx in rename_contexts:
        edge_stats = self._rebuild_connect_edges(path, ctx)
        self._merge_stats(total_stats, edge_stats)

    for path, ctx in mod_contexts:
        edge_stats = self._incremental_connect_edges(path, ctx)
        self._merge_stats(total_stats, edge_stats)

    nodes_after = self.code_graph.graph.vcount()
    edges_after = self.code_graph.graph.ecount()
    total_stats["nodes_before"] = nodes_before
    total_stats["nodes_after"] = nodes_after
    total_stats["edges_before"] = edges_before
    total_stats["edges_after"] = edges_after

    commit_info = ""
    if earlier_commit:
        total_stats["commit_earlier"] = earlier_commit
        total_stats["commit_later"] = effective_later_commit
        commit_info = (
            f"commits {earlier_commit[:12]}..{effective_later_commit[:12]}, "
        )

    total_changed = sum(
        len(changed_files.get(k, []))
        for k in ("deleted", "added", "renamed", "modified")
    )
    logger.info(
        f"Incremental patch summary: {commit_info}"
        f"{total_changed} files changed, "
        f"nodes {nodes_before}{nodes_after}{nodes_after - nodes_before:+d}), "
        f"edges {edges_before}{edges_after}{edges_after - edges_before:+d})"
    )

    # Rebuild line-range indexes after each patcher batch so range
    # queries reflect the post-patch graph state. Cost is O(V+E),
    # negligible relative to the patch itself.
    with self.profiler.section("patch_files.build_range_indexes"):
        self.code_graph.build_range_indexes()

    # Profiler summary — also stash structured timings on the patcher so
    # benchmarks can read them without parsing logs.
    report = self.profiler.report(reset=True)
    total_stats["profile"] = {
        label: {"total_s": st.total, "count": st.count} for label, st in report
    }

    return total_stats

MultiGraphIndex

MultiGraphIndex(
    code_graph: CodeGraph,
    specs: Mapping[str, GraphLayerSpec] | None = None,
    *,
    use_core: bool = True
)

Layer index for a :class:CodeGraph.

The index stores edge-id lists per layer. Layers can overlap; today every reference edge also belongs to dependency and all.

Source code in codenib/graph/layers.py
def __init__(
    self,
    code_graph: CodeGraph,
    specs: Optional[Mapping[str, GraphLayerSpec]] = None,
    *,
    use_core: bool = True,
):
    self.code_graph = code_graph
    self.specs = _normalize_specs(specs or DEFAULT_LAYER_SPECS)
    self._edge_ids_by_layer = classify_edge_layers(
        _edge_types(code_graph),
        self.specs,
        use_core=use_core,
    )

ROISubgraph

ROISubgraph(code_graph: CodeGraph)

Region of Interest (ROI) subgraph extraction for code graphs.

This class takes search results from CodeSearchEngine and extracts a focused subgraph containing the most relevant code elements and their neighborhood.

Parameters:

Name Type Description Default
code_graph CodeGraph

The full code graph from which to extract subgraphs

required

Methods:

Name Description
expand_ppr

Expand seed nodes using Personalized PageRank.

extract_subgraph

Extract a subgraph around specified node names.

get_filtered_subgraph_nodes

Extract useful nodes from a subgraph, filtering out nodes where

get_node_info_by_id

Get information about a node in the original graph.

get_node_content

Get the content of a node in the original graph.

Source code in codenib/graph/roi_subgraph.py
def __init__(self, code_graph: CodeGraph):
    """
    Initialize the ROI subgraph extractor.

    Args:
        code_graph: The full code graph from which to extract subgraphs
    """
    self.code_graph = code_graph
    self.full_graph = code_graph.get_graph()

expand_ppr

expand_ppr(
    node_names: list[str],
    top_k: int = 50,
    damping: float = 0.85,
    filter_tests: bool = True,
) -> list[NodeInfo]

Expand seed nodes using Personalized PageRank.

Runs PPR on the full graph with seed nodes as the personalization vector, then returns the top-k nodes ranked by PPR score.

Parameters:

Name Type Description Default
node_names list[str]

Seed node names (e.g. from BM25).

required
top_k int

Maximum number of nodes to return.

50
damping float

PPR damping factor (0–1). Higher = more global.

0.85
filter_tests bool

Whether to exclude test files.

True

Returns:

Type Description
list[NodeInfo]

List of NodeInfo objects sorted by PPR score (descending).

Source code in codenib/graph/roi_subgraph.py
def expand_ppr(
    self,
    node_names: List[str],
    top_k: int = 50,
    damping: float = 0.85,
    filter_tests: bool = True,
) -> List[NodeInfo]:
    """Expand seed nodes using Personalized PageRank.

    Runs PPR on the full graph with seed nodes as the personalization
    vector, then returns the top-k nodes ranked by PPR score.

    Args:
        node_names: Seed node names (e.g. from BM25).
        top_k: Maximum number of nodes to return.
        damping: PPR damping factor (0–1). Higher = more global.
        filter_tests: Whether to exclude test files.

    Returns:
        List of NodeInfo objects sorted by PPR score (descending).
    """
    n_vertices = self.full_graph.vcount()
    if n_vertices == 0:
        return []

    # Build personalization vector: uniform weight on seeds, 0 elsewhere
    reset = [0.0] * n_vertices
    seed_count = 0
    for name in node_names:
        vid = self.code_graph.name_to_vertex.get(name)
        if vid is not None:
            reset[vid] = 1.0
            seed_count += 1
        else:
            logger.warning("PPR seed '%s' not found in graph", name)

    if seed_count == 0:
        logger.warning("No valid PPR seeds — returning empty results")
        return []

    # Run Personalized PageRank
    scores = self.full_graph.personalized_pagerank(
        directed=True, damping=damping, reset=reset
    )

    # Pair each vertex with its PPR score and sort descending
    scored = sorted(enumerate(scores), key=lambda x: x[1], reverse=True)

    results: List[NodeInfo] = []
    for vid, ppr_score in scored:
        if len(results) >= top_k:
            break

        node_attrs = self.get_node_info_by_id(vid)

        # Filter test files
        if filter_tests and node_attrs.file and is_test_file(node_attrs.file):
            continue

        # Skip zero-span or missing line info
        if (
            node_attrs.start_line is None
            or node_attrs.end_line is None
            or node_attrs.start_line == node_attrs.end_line
        ):
            continue

        content = self.get_node_content(vid) or ""
        results.append(
            NodeInfo(
                node_name=node_attrs.node_name,
                type=node_attrs.type,
                file=node_attrs.file,
                start_line=node_attrs.start_line,
                end_line=node_attrs.end_line,
                score=ppr_score,
                content=content,
            )
        )

    logger.info(
        "PPR expansion: %d seeds -> %d nodes (damping=%.2f)",
        seed_count,
        len(results),
        damping,
    )
    return results

extract_subgraph

extract_subgraph(
    node_names: list[str],
    k_hop: int = 2,
    edge_types: list[str] | None = None,
    direction: str = "both",
) -> Graph

Extract a subgraph around specified node names.

Parameters:

Name Type Description Default
node_names list[str]

List of node names to use as seeds for the subgraph

required
k_hop int

Number of hops to expand from each seed node

2
edge_types list[str] | None

Optional list of edge types to consider (None for all types)

None
direction str

Direction to traverse - "forward", "backward", or "both"

'both'

Returns:

Type Description
Graph

An igraph Graph object representing the extracted subgraph

Source code in codenib/graph/roi_subgraph.py
def extract_subgraph(
    self,
    node_names: List[str],
    k_hop: int = 2,
    edge_types: Optional[List[str]] = None,
    direction: str = "both",
) -> ig.Graph:
    """
    Extract a subgraph around specified node names.

    Args:
        node_names: List of node names to use as seeds for the subgraph
        k_hop: Number of hops to expand from each seed node
        edge_types: Optional list of edge types to consider (None for all types)
        direction: Direction to traverse - "forward", "backward", or "both"

    Returns:
        An igraph Graph object representing the extracted subgraph
    """
    # Track visited nodes to avoid duplicates
    visited_nodes = set()
    # Track nodes to include in the subgraph
    subgraph_nodes = set()

    # BFS for each seed node
    for node_name in node_names:
        try:
            # Convert node name to node ID
            if node_name not in self.code_graph.name_to_vertex:
                logger.warning(f"Node name {node_name!r} not found in graph")
                continue

            node_id = self.code_graph.name_to_vertex[node_name]

            # Skip if we've already processed this node
            if node_id in visited_nodes:
                continue

            # Add this seed node to the subgraph and mark as visited
            subgraph_nodes.add(node_id)
            visited_nodes.add(node_id)

            # Initialize queue for BFS with (node, depth) pairs
            queue = deque([(node_id, 0)])

            # BFS traversal to find k-hop neighborhood
            while queue:
                current_node, depth = queue.popleft()

                # Stop expanding if we've reached k hops
                if depth >= k_hop:
                    continue

                # Get neighbors
                neighbors = self._get_neighbors(current_node, edge_types, direction)

                # Process neighbors
                for neighbor in neighbors:
                    if neighbor not in visited_nodes:
                        visited_nodes.add(neighbor)
                        subgraph_nodes.add(neighbor)
                        queue.append((neighbor, depth + 1))

        except Exception as e:
            logger.error(f"Error processing node {node_name}: {e}")
            continue

    # Create subgraph from the collected nodes
    return self._create_subgraph(subgraph_nodes)

get_filtered_subgraph_nodes

get_filtered_subgraph_nodes(
    subgraph: Graph,
    exclude_nodes: list[str] | None = None,
    filter_tests: bool = True,
    node_types: list[str] | None = None,
) -> list[NodeInfo]

Extract useful nodes from a subgraph, filtering out nodes where start_line equals end_line (unless explicitly included).

Parameters:

Name Type Description Default
subgraph Graph

An igraph Graph object (from extract_subgraph or extract_roi_from_search_results)

required
node_types list[str] | None

Optional list of node types to include (None for all types)

None

Returns:

Type Description
list[NodeInfo]

List of NodeInfo objects including the node content

Source code in codenib/graph/roi_subgraph.py
def get_filtered_subgraph_nodes(
    self,
    subgraph: ig.Graph,
    exclude_nodes: Optional[List[str]] = None,
    filter_tests: bool = True,
    node_types: Optional[List[str]] = None,
) -> List[NodeInfo]:
    """
    Extract useful nodes from a subgraph, filtering out nodes where
    start_line equals end_line (unless explicitly included).

    Args:
        subgraph: An igraph Graph object (from extract_subgraph or
            extract_roi_from_search_results)
        node_types: Optional list of node types to include (None for all types)

    Returns:
        List of NodeInfo objects including the node content
    """
    filtered_nodes = []

    for node in subgraph.vs:
        try:
            # Get original node ID and attributes
            original_id = node["original_id"]
            node_attrs = self.get_node_info_by_id(original_id)

            # Filter out test files if specified
            if filter_tests and node_attrs.file and is_test_file(node_attrs.file):
                continue

            # Filter by node type if specified
            if node_types and node_attrs.type not in node_types:
                continue

            # Skip nodes with zero line span or None values
            if (
                node_attrs.start_line is None
                or node_attrs.end_line is None
                or node_attrs.start_line == node_attrs.end_line
            ):
                continue

            # Get the node content
            content = self.get_node_content(original_id) or ""

            # Create NodeInfo and add to the list
            node_with_content = NodeInfo(
                node_name=node_attrs.node_name,
                type=node_attrs.type,
                file=node_attrs.file,
                start_line=node_attrs.start_line,
                end_line=node_attrs.end_line,
                content=content,
            )

            filtered_nodes.append(node_with_content)

        except Exception as e:
            logger.error(f"Error processing node {node.index}: {e}")
            continue

    # remove exclude_nodex in filtered_nodes
    if exclude_nodes:
        filtered_nodes = [
            node for node in filtered_nodes if node.node_name not in exclude_nodes
        ]

    logger.info(f"Extracted {len(filtered_nodes)} useful nodes from subgraph")
    return filtered_nodes

get_node_info_by_id

get_node_info_by_id(node_id: int) -> NodeInfo

Get information about a node in the original graph.

Parameters:

Name Type Description Default
node_id int

ID of the node in the original graph

required

Returns:

Type Description
NodeInfo

NodeInfo containing the node attributes

Source code in codenib/graph/roi_subgraph.py
def get_node_info_by_id(self, node_id: int) -> NodeInfo:
    """
    Get information about a node in the original graph.

    Args:
        node_id: ID of the node in the original graph

    Returns:
        NodeInfo containing the node attributes
    """
    try:
        vertex = self.full_graph.vs[node_id]
        attributes = vertex.attributes()

        # In igraph, the vertex name is stored in the "name" attribute
        node_name = vertex["name"]

        node_info_dict = {
            "node_name": node_name,
            "type": attributes.get("type", ""),
            "file": attributes.get("file", None),
            "start_line": attributes.get("start_line", None),
            "end_line": attributes.get("end_line", None),
        }

        return NodeInfo(**node_info_dict)
    except Exception as e:
        logger.error(f"Error getting node info for node {node_id}: {e}")
        return NodeInfo(type="")

get_node_content

get_node_content(node_id: int) -> str | None

Get the content of a node in the original graph.

Parameters:

Name Type Description Default
node_id int

ID of the node in the original graph

required

Returns: The content of the node, or None if not available

Source code in codenib/graph/roi_subgraph.py
def get_node_content(self, node_id: int) -> Optional[str]:
    """
    Get the content of a node in the original graph.

    Args:
        node_id: ID of the node in the original graph
    Returns:
        The content of the node, or None if not available
    """
    node_content = self.code_graph.get_node_content(node_id)
    if node_content is None:
        logger.warning(f"No content found for node {node_id}")
        return None
    return node_content

build_graph_layers

build_graph_layers(
    code_graph: CodeGraph,
    specs: Mapping[str, GraphLayerSpec] | None = None,
    *,
    use_core: bool = True
) -> MultiGraphIndex

Build a layer index for code_graph.

Source code in codenib/graph/layers.py
def build_graph_layers(
    code_graph: CodeGraph,
    specs: Optional[Mapping[str, GraphLayerSpec]] = None,
    *,
    use_core: bool = True,
) -> MultiGraphIndex:
    """Build a layer index for ``code_graph``."""
    return MultiGraphIndex(code_graph, specs, use_core=use_core)

traverse_tree_structure

traverse_tree_structure(
    code_graph: CodeGraph,
    root,
    direction="downstream",
    hops=2,
    node_type_filter: list[str] | None = None,
    edge_type_filter: list[str] | None = None,
)

Traverse tree structure starting from a root node

Parameters:

Name Type Description Default
code_graph CodeGraph

CodeGraph instance

required
root

Root node ID to start traversal from

required
direction

'downstream', 'upstream', or 'both'

'downstream'
hops

Maximum number of hops to traverse (-1 for unlimited)

2
node_type_filter list[str] | None

Filter by node types

None
edge_type_filter list[str] | None

Filter by edge types

None

Returns:

Type Description

String representation of the tree structure

Source code in codenib/graph/traverse_graph.py
def traverse_tree_structure(
    code_graph: CodeGraph,
    root,
    direction="downstream",
    hops=2,
    node_type_filter: Optional[List[str]] = None,
    edge_type_filter: Optional[List[str]] = None,
):
    """
    Traverse tree structure starting from a root node

    Args:
        code_graph: CodeGraph instance
        root: Root node ID to start traversal from
        direction: 'downstream', 'upstream', or 'both'
        hops: Maximum number of hops to traverse (-1 for unlimited)
        node_type_filter: Filter by node types
        edge_type_filter: Filter by edge types

    Returns:
        String representation of the tree structure
    """
    if hops == -1:
        hops = 20

    if root not in code_graph.name_to_vertex:
        return f"Node {root!r} not found in graph"

    rtn_str = []
    traversed_nodes = set()
    traversed_edges = set()

    def traverse(node, prefix, is_last, level, edge_type, edirection):
        if level > hops:
            return

        if node == root and level == 0:
            rtn_str.append(f"{node}")
            new_prefix = ""
            edirection = direction
        else:
            connector = "└── " if is_last else "├── "
            connector += f"{edge_type} ── "
            rtn_str.append(f"{prefix}{connector}{node}")
            new_prefix = prefix + (" " if is_last else "│") + " " * (len(connector) - 1)

        if node in traversed_nodes:
            return
        traversed_nodes.add(node)

        # Separate containment and reference edges
        contain_neighbors = []  # (neighbor_id, etype, edir)
        reference_neighbors = []  # (neighbor_id, etype, edir)

        def is_ntype_not_valid(_ntype):
            return node_type_filter is not None and _ntype not in node_type_filter

        def is_etype_not_valid(_etype):
            return edge_type_filter is not None and _etype not in edge_type_filter

        if node not in code_graph.name_to_vertex:
            return

        vertex_id = code_graph.name_to_vertex[node]

        # Downstream traversal — iterate edges directly so multi-edges
        # between the same pair are each surfaced with their own type.
        if "downstream" == edirection or (node == root and direction == "both"):
            for eid in code_graph.graph.incident(vertex_id, mode="out"):
                edge = code_graph.graph.es[eid]
                neighbor_id = edge.target
                neighbor_vertex = code_graph.graph.vs[neighbor_id]
                neighbor = neighbor_vertex["name"]
                neigh_type = (
                    neighbor_vertex["type"]
                    if "type" in neighbor_vertex.attributes()
                    else "unknown"
                )

                if is_ntype_not_valid(neigh_type):
                    continue

                etype = edge["type"] if "type" in edge.attributes() else "unknown"
                if is_etype_not_valid(etype):
                    continue
                if is_test_file(neighbor):
                    continue
                if (node, etype, neighbor) in traversed_edges:
                    continue
                traversed_edges.add((node, etype, neighbor))
                if etype == "contain":
                    contain_neighbors.append((neighbor, etype, "downstream"))
                else:
                    reference_neighbors.append((neighbor, etype, "downstream"))

        # Upstream traversal
        if "upstream" == edirection or (node == root and direction == "both"):
            for eid in code_graph.graph.incident(vertex_id, mode="in"):
                edge = code_graph.graph.es[eid]
                neighbor_id = edge.source
                neighbor_vertex = code_graph.graph.vs[neighbor_id]
                neighbor = neighbor_vertex["name"]
                neigh_type = (
                    neighbor_vertex["type"]
                    if "type" in neighbor_vertex.attributes()
                    else "unknown"
                )

                if is_ntype_not_valid(neigh_type):
                    continue

                etype = edge["type"] if "type" in edge.attributes() else "unknown"
                if is_etype_not_valid(etype):
                    continue
                if is_test_file(neighbor):
                    continue
                if (neighbor, etype, node) in traversed_edges:
                    continue
                traversed_edges.add((neighbor, etype, node))
                if etype == "contain":
                    contain_neighbors.append((neighbor, etype, "upstream"))
                else:
                    reference_neighbors.append((neighbor, etype, "upstream"))

        # Combine: containment first, then references
        all_neighbors = contain_neighbors + reference_neighbors

        for i, (neigh_id, etype, edir) in enumerate(all_neighbors):
            is_last_child = i == len(all_neighbors) - 1
            if edir == "upstream":
                etype += "-by"
            traverse(neigh_id, new_prefix, is_last_child, level + 1, etype, edir)

    traverse(root, "", False, 0, None, None)
    return "\n".join(rtn_str)