mirror of
https://github.com/safishamsi/graphify.git
synced 2026-07-14 11:27:10 +00:00
Safi
cce2673021
llm.py: add explicit edge direction rule to extraction system prompt — source = actor (caller/importer/subclass), target = acted-upon (callee/ imported/base). LLM was systematically emitting callee->caller for calls edges because the schema never stated direction semantics. build.py: extend ghost-node merge to catch LLM nodes that populate source_location (bypassing the old None check). Now uses _origin=="ast" as the canonical signal — AST nodes always win; any non-AST node sharing (basename, label) with an AST node is collapsed into the AST canonical. Fixes LLM bare-stem IDs (bpe_get_pairs) surviving alongside AST parent-qualified IDs (mingpt_bpe_get_pairs) and carrying reversed edges. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
500 lines
22 KiB
Python
500 lines
22 KiB
Python
# assemble node+edge dicts into a NetworkX graph, preserving edge direction
|
||
#
|
||
# Node deduplication — three layers:
|
||
#
|
||
# 1. Within a file (AST): each extractor tracks a `seen_ids` set. A node ID is
|
||
# emitted at most once per file, so duplicate class/function definitions in
|
||
# the same source file are collapsed to the first occurrence.
|
||
#
|
||
# 2. Between files (build): NetworkX G.add_node() is idempotent — calling it
|
||
# twice with the same ID overwrites the attributes with the second call's
|
||
# values. Nodes are added in extraction order (AST first, then semantic),
|
||
# so if the same entity is extracted by both passes the semantic node
|
||
# silently overwrites the AST node. This is intentional: semantic nodes
|
||
# carry richer labels and cross-file context, while AST nodes have precise
|
||
# source_location. If you need to change the priority, reorder extractions
|
||
# passed to build().
|
||
#
|
||
# 3. Semantic merge (skill): before calling build(), the skill merges cached
|
||
# and new semantic results using an explicit `seen` set keyed on node["id"],
|
||
# so duplicates across cache hits and new extractions are resolved there
|
||
# before any graph construction happens.
|
||
#
|
||
from __future__ import annotations
|
||
import json
|
||
import os
|
||
import re
|
||
import sys
|
||
import unicodedata
|
||
from pathlib import Path
|
||
import networkx as nx
|
||
from .validate import validate_extraction
|
||
|
||
|
||
# Synonym mapper for known invalid file_type values that LLM subagents commonly
|
||
# emit. Keeps semantic intent close (markdown→document, tool→code) and falls
|
||
# back to "concept" for any other invalid value (see #840).
|
||
_FILE_TYPE_SYNONYMS = {
|
||
"markdown": "document",
|
||
"text": "document",
|
||
"tool": "code",
|
||
"library": "code",
|
||
"pattern": "concept",
|
||
"principle": "concept",
|
||
"constraint": "concept",
|
||
"tech": "concept",
|
||
"technology": "concept",
|
||
"data-source": "concept",
|
||
"data_source": "concept",
|
||
"gotcha": "concept",
|
||
"framework": "concept",
|
||
}
|
||
|
||
|
||
def _normalize_id(s: str) -> str:
|
||
r"""Normalize an ID string the same way extract._make_id does.
|
||
|
||
Used to reconcile edge endpoints when the LLM generates IDs with slightly
|
||
different punctuation or casing than the AST extractor. Must stay in sync
|
||
with extract._make_id — NFKC normalization, \w with re.UNICODE, underscore
|
||
collapse, and casefold must all match (#811).
|
||
"""
|
||
s = unicodedata.normalize("NFKC", s)
|
||
cleaned = re.sub(r"[^\w]+", "_", s, flags=re.UNICODE)
|
||
cleaned = re.sub(r"_+", "_", cleaned)
|
||
return cleaned.strip("_").casefold()
|
||
|
||
|
||
def _norm_source_file(p: str | None, root: str | None = None) -> str | None:
|
||
"""Normalize path separators and relativize absolute paths.
|
||
|
||
Converts backslashes to forward slashes (Windows compatibility) and, when
|
||
root is provided, strips the absolute prefix from paths produced by semantic
|
||
subagents so source_file is always repo-relative (fixes #932).
|
||
"""
|
||
if not p:
|
||
return p
|
||
p = p.replace("\\", "/")
|
||
if root and os.path.isabs(p):
|
||
try:
|
||
p = Path(p).relative_to(root).as_posix()
|
||
except ValueError:
|
||
pass
|
||
return p
|
||
|
||
|
||
def edge_data(G: nx.Graph, u: str, v: str) -> dict:
|
||
"""Return one edge attribute dict for (u, v), tolerating MultiGraph.
|
||
|
||
For MultiGraph/MultiDiGraph there can be multiple parallel edges;
|
||
this returns the first one (sufficient for callers that only need
|
||
relation/confidence for rendering). Fixes #796.
|
||
"""
|
||
raw = G[u][v]
|
||
if isinstance(G, (nx.MultiGraph, nx.MultiDiGraph)):
|
||
return next(iter(raw.values()), {})
|
||
return raw
|
||
|
||
|
||
def edge_datas(G: nx.Graph, u: str, v: str) -> list[dict]:
|
||
"""Return every edge attribute dict for (u, v); always a list."""
|
||
raw = G[u][v]
|
||
if isinstance(G, (nx.MultiGraph, nx.MultiDiGraph)):
|
||
return list(raw.values())
|
||
return [raw]
|
||
|
||
|
||
def build_from_json(extraction: dict, *, directed: bool = False, root: str | Path | None = None) -> nx.Graph:
|
||
"""Build a NetworkX graph from an extraction dict.
|
||
|
||
directed=True produces a DiGraph that preserves edge direction (source→target).
|
||
directed=False (default) produces an undirected Graph for backward compatibility.
|
||
root: if given, absolute source_file paths from semantic subagents are made
|
||
relative to root so all nodes share a consistent path key (#932).
|
||
"""
|
||
_root = str(Path(root).resolve()) if root else None
|
||
# NetworkX <= 3.1 serialised edges as "links"; remap to "edges" for compatibility.
|
||
if "edges" not in extraction and "links" in extraction:
|
||
extraction = dict(extraction, edges=extraction["links"])
|
||
|
||
# Canonicalize legacy node/edge schema before validation.
|
||
for node in extraction.get("nodes", []):
|
||
if not isinstance(node, dict):
|
||
continue
|
||
if "source" in node and "source_file" not in node:
|
||
# Count edges that reference this node so the warning is actionable (#479)
|
||
node_id = node.get("id", "?")
|
||
affected_edges = sum(
|
||
1 for e in extraction.get("edges", [])
|
||
if e.get("source") == node_id or e.get("target") == node_id
|
||
)
|
||
print(
|
||
f"[graphify] WARNING: node '{node_id}' uses field 'source' instead of "
|
||
f"'source_file' — {affected_edges} edge(s) may be misrouted. "
|
||
f"Rename the field to 'source_file' to silence this warning.",
|
||
file=sys.stderr,
|
||
)
|
||
node["source_file"] = node.pop("source")
|
||
# Default missing/None file_type to "concept" so legacy graph.json
|
||
# entries (and stub nodes preserved by `_rebuild_code` from older
|
||
# graphify versions that didn't always populate file_type) don't
|
||
# trigger spurious "invalid file_type 'None'" validator warnings (#660).
|
||
if node.get("file_type") in (None, ""):
|
||
node["file_type"] = "concept"
|
||
ft = node.get("file_type", "")
|
||
if ft and ft not in {"code", "document", "paper", "image", "rationale", "concept"}:
|
||
node["file_type"] = _FILE_TYPE_SYNONYMS.get(ft, "concept")
|
||
|
||
errors = validate_extraction(extraction)
|
||
# Dangling edges (stdlib/external imports) are expected - only warn about real schema errors.
|
||
real_errors = [e for e in errors if "does not match any node id" not in e]
|
||
if real_errors:
|
||
print(f"[graphify] Extraction warning ({len(real_errors)} issues): {real_errors[0]}", file=sys.stderr)
|
||
G: nx.Graph = nx.DiGraph() if directed else nx.Graph()
|
||
for node in extraction.get("nodes", []):
|
||
if "source_file" in node:
|
||
node["source_file"] = _norm_source_file(node["source_file"], _root)
|
||
G.add_node(node["id"], **{k: v for k, v in node.items() if k != "id"})
|
||
node_set = set(G.nodes())
|
||
|
||
# #1145 (extended): merge LLM ghost-duplicate nodes into AST canonical nodes.
|
||
# Original bug: AST uses parent-qualified IDs (mingpt_bpe_get_pairs) while LLM
|
||
# uses bare-stem IDs (bpe_get_pairs) — different IDs, same symbol.
|
||
# Original fix only caught LLM nodes with source_location=None; LLM now
|
||
# populates source_location, so those ghosts survived. Extended fix: use
|
||
# _origin=="ast" as the canonical signal. AST nodes always win; any non-AST
|
||
# node sharing (basename, label) with an AST node is a ghost.
|
||
_loc_nodes: dict[tuple[str, str], str] = {} # (basename, label) -> canonical node id
|
||
_noloc_nodes: dict[tuple[str, str], str] = {} # (basename, label) -> ghost node id
|
||
|
||
# Pass 1: collect canonical nodes — AST-origin nodes take precedence over LLM nodes.
|
||
for nid in node_set:
|
||
attrs = G.nodes[nid]
|
||
label = str(attrs.get("label", "")).strip()
|
||
sf = str(attrs.get("source_file", ""))
|
||
basename = Path(sf).name if sf else ""
|
||
if not label or not basename:
|
||
continue
|
||
if attrs.get("source_location") or attrs.get("_origin") == "ast":
|
||
key = (basename, label)
|
||
# AST-origin nodes always overwrite; non-AST only written if key unseen.
|
||
if attrs.get("_origin") == "ast" or key not in _loc_nodes:
|
||
_loc_nodes[key] = nid
|
||
|
||
# Pass 2: find ghosts — non-AST nodes that have an AST canonical twin.
|
||
for nid in node_set:
|
||
attrs = G.nodes[nid]
|
||
if attrs.get("_origin") == "ast":
|
||
continue # AST nodes are never ghosts
|
||
label = str(attrs.get("label", "")).strip()
|
||
sf = str(attrs.get("source_file", ""))
|
||
basename = Path(sf).name if sf else ""
|
||
if not label or not basename:
|
||
continue
|
||
key = (basename, label)
|
||
if key in _loc_nodes and _loc_nodes[key] != nid:
|
||
_noloc_nodes[key] = nid
|
||
# For every ghost that has an AST counterpart, record a remap.
|
||
_ghost_remap: dict[str, str] = {} # ghost_id -> canonical_id
|
||
for key, sem_id in _noloc_nodes.items():
|
||
ast_id = _loc_nodes.get(key)
|
||
if ast_id is not None:
|
||
_ghost_remap[sem_id] = ast_id
|
||
# Remove ghost nodes from the graph; edges will be re-pointed via norm_to_id.
|
||
for ghost_id in _ghost_remap:
|
||
G.remove_node(ghost_id)
|
||
node_set.discard(ghost_id)
|
||
|
||
# Normalized ID map: lets edges survive when the LLM generates IDs with
|
||
# slightly different casing or punctuation than the AST extractor.
|
||
# e.g. "Session_ValidateToken" maps to "session_validatetoken".
|
||
norm_to_id: dict[str, str] = {_normalize_id(nid): nid for nid in node_set}
|
||
# Also map ghost IDs to their canonical AST replacements.
|
||
for ghost_id, canonical_id in _ghost_remap.items():
|
||
norm_to_id[_normalize_id(ghost_id)] = canonical_id
|
||
norm_to_id[ghost_id] = canonical_id
|
||
# Iterate edges in a deterministic order. The graph is undirected and stores
|
||
# direction in _src/_tgt; when two edges collapse onto the same node pair the
|
||
# last write wins, so an unstable iteration order flips _src/_tgt run-to-run
|
||
# and makes the serialized graph churn. Sorting fixes the last-write outcome.
|
||
for edge in sorted(
|
||
extraction.get("edges", []),
|
||
key=lambda e: (
|
||
str(e.get("source", e.get("from", ""))),
|
||
str(e.get("target", e.get("to", ""))),
|
||
str(e.get("relation", "")),
|
||
),
|
||
):
|
||
if "source" not in edge and "from" in edge:
|
||
edge["source"] = edge["from"]
|
||
if "target" not in edge and "to" in edge:
|
||
edge["target"] = edge["to"]
|
||
if "source" not in edge or "target" not in edge:
|
||
continue
|
||
src, tgt = edge["source"], edge["target"]
|
||
# Remap mismatched IDs via normalization before dropping the edge.
|
||
if src not in node_set:
|
||
src = norm_to_id.get(_normalize_id(src), src)
|
||
if tgt not in node_set:
|
||
tgt = norm_to_id.get(_normalize_id(tgt), tgt)
|
||
if src not in node_set or tgt not in node_set:
|
||
continue # skip edges to external/stdlib nodes - expected, not an error
|
||
attrs = {k: v for k, v in edge.items() if k not in ("source", "target")}
|
||
if "source_file" in attrs:
|
||
attrs["source_file"] = _norm_source_file(attrs["source_file"], _root)
|
||
# Drop cross-language INFERRED `calls` edges — same short names (render,
|
||
# parse, etc.) appear across language boundaries in multi-language chunks,
|
||
# producing phantom edges that don't represent real call relationships.
|
||
if attrs.get("relation") == "calls" and attrs.get("confidence") == "INFERRED":
|
||
_LANG_FAMILY: dict[str, str] = {
|
||
".py": "py", ".pyi": "py",
|
||
".js": "js", ".mjs": "js", ".cjs": "js", ".jsx": "js",
|
||
".ts": "js", ".tsx": "js",
|
||
".go": "go", ".rs": "rs",
|
||
".java": "jvm", ".kt": "jvm", ".scala": "jvm", ".groovy": "jvm",
|
||
".c": "c", ".h": "c", ".cc": "cpp", ".cpp": "cpp", ".hpp": "cpp",
|
||
".rb": "rb", ".php": "php", ".cs": "cs", ".swift": "swift", ".lua": "lua",
|
||
}
|
||
src_ext = Path(G.nodes[src].get("source_file") or "").suffix.lower()
|
||
tgt_ext = Path(G.nodes[tgt].get("source_file") or "").suffix.lower()
|
||
if src_ext and tgt_ext and _LANG_FAMILY.get(src_ext) != _LANG_FAMILY.get(tgt_ext):
|
||
continue
|
||
# Preserve original edge direction - undirected graphs lose it otherwise,
|
||
# causing display functions to show edges backwards.
|
||
attrs["_src"] = src
|
||
attrs["_tgt"] = tgt
|
||
# When the graph is undirected and the same node pair appears twice with
|
||
# the same relation but opposite directions (e.g. a `calls` b and b `calls` a),
|
||
# nx.Graph collapses them into one edge. The deterministic sort above means
|
||
# the lexicographically-later direction would systematically overwrite the
|
||
# earlier one's _src/_tgt, silently flipping the surviving edge's caller
|
||
# and callee. First-seen direction wins instead — drop the redundant
|
||
# reverse-direction duplicate so the original direction is preserved (#1061).
|
||
if not G.is_directed() and G.has_edge(src, tgt):
|
||
existing = edge_data(G, src, tgt)
|
||
if existing.get("relation") == attrs.get("relation") and (
|
||
existing.get("_src") == tgt and existing.get("_tgt") == src
|
||
):
|
||
continue
|
||
G.add_edge(src, tgt, **attrs)
|
||
hyperedges = extraction.get("hyperedges", [])
|
||
if hyperedges:
|
||
G.graph["hyperedges"] = hyperedges
|
||
return G
|
||
|
||
|
||
def build(
|
||
extractions: list[dict],
|
||
*,
|
||
directed: bool = False,
|
||
dedup: bool = True,
|
||
dedup_llm_backend: str | None = None,
|
||
root: str | Path | None = None,
|
||
) -> nx.Graph:
|
||
"""Merge multiple extraction results into one graph.
|
||
|
||
directed=True produces a DiGraph that preserves edge direction (source→target).
|
||
directed=False (default) produces an undirected Graph for backward compatibility.
|
||
dedup=True (default) runs entity deduplication before building the graph.
|
||
dedup_llm_backend: if set (e.g. "gemini", "claude", or "kimi"), uses LLM to resolve
|
||
ambiguous pairs in the 75–92 Jaro-Winkler score zone.
|
||
root: if given, absolute source_file paths are made relative to root (#932).
|
||
|
||
Extractions are merged in order. For nodes with the same ID, the last
|
||
extraction's attributes win (NetworkX add_node overwrites). Pass AST
|
||
results before semantic results so semantic labels take precedence, or
|
||
reverse the order if you prefer AST source_location precision to win.
|
||
"""
|
||
from graphify.dedup import deduplicate_entities
|
||
combined: dict = {"nodes": [], "edges": [], "hyperedges": [], "input_tokens": 0, "output_tokens": 0}
|
||
for ext in extractions:
|
||
combined["nodes"].extend(ext.get("nodes", []))
|
||
combined["edges"].extend(ext.get("edges", []))
|
||
combined["hyperedges"].extend(ext.get("hyperedges", []))
|
||
combined["input_tokens"] += ext.get("input_tokens", 0)
|
||
combined["output_tokens"] += ext.get("output_tokens", 0)
|
||
if dedup and combined["nodes"]:
|
||
combined["nodes"], combined["edges"] = deduplicate_entities(
|
||
combined["nodes"], combined["edges"], communities={},
|
||
dedup_llm_backend=dedup_llm_backend,
|
||
)
|
||
return build_from_json(combined, directed=directed, root=root)
|
||
|
||
|
||
def _norm_label(label: str | None) -> str:
|
||
"""Canonical dedup key — Unicode-aware, preserves CJK/word characters."""
|
||
if not isinstance(label, str):
|
||
label = "" if label is None else str(label)
|
||
label = unicodedata.normalize("NFKC", label)
|
||
return re.sub(r"[\W_ ]+", " ", label.casefold(), flags=re.UNICODE).strip()
|
||
|
||
|
||
def deduplicate_by_label(nodes: list[dict], edges: list[dict]) -> tuple[list[dict], list[dict]]:
|
||
"""Merge nodes that share a normalised label, rewriting edge references.
|
||
|
||
Prefers IDs without chunk suffixes (_c\\d+) and shorter IDs when tied.
|
||
Drops self-loops created by the merge. Called in build() automatically.
|
||
"""
|
||
_CHUNK_SUFFIX = re.compile(r"_c\d+$")
|
||
canonical: dict[str, dict] = {} # norm_label -> surviving node
|
||
remap: dict[str, str] = {} # old_id -> surviving_id
|
||
|
||
for node in nodes:
|
||
key = _norm_label(node.get("label", node.get("id", "")))
|
||
if not key:
|
||
continue
|
||
existing = canonical.get(key)
|
||
if existing is None:
|
||
canonical[key] = node
|
||
else:
|
||
has_suffix = bool(_CHUNK_SUFFIX.search(node["id"]))
|
||
existing_has_suffix = bool(_CHUNK_SUFFIX.search(existing["id"]))
|
||
if has_suffix and not existing_has_suffix:
|
||
remap[node["id"]] = existing["id"]
|
||
elif existing_has_suffix and not has_suffix:
|
||
remap[existing["id"]] = node["id"]
|
||
canonical[key] = node
|
||
elif len(node["id"]) < len(existing["id"]):
|
||
remap[existing["id"]] = node["id"]
|
||
canonical[key] = node
|
||
else:
|
||
remap[node["id"]] = existing["id"]
|
||
|
||
if not remap:
|
||
return nodes, edges
|
||
|
||
print(f"[graphify] Deduplicated {len(remap)} duplicate node(s) by label.", file=sys.stderr)
|
||
deduped_nodes = list(canonical.values())
|
||
deduped_edges = []
|
||
for edge in edges:
|
||
e = dict(edge)
|
||
e["source"] = remap.get(e["source"], e["source"])
|
||
e["target"] = remap.get(e["target"], e["target"])
|
||
if e["source"] != e["target"]:
|
||
deduped_edges.append(e)
|
||
return deduped_nodes, deduped_edges
|
||
|
||
|
||
def build_merge(
|
||
new_chunks: list[dict],
|
||
graph_path: str | Path = "graphify-out/graph.json",
|
||
prune_sources: list[str] | None = None,
|
||
*,
|
||
directed: bool = False,
|
||
dedup: bool = True,
|
||
dedup_llm_backend: str | None = None,
|
||
root: str | Path | None = None,
|
||
) -> nx.Graph:
|
||
"""Load existing graph.json, merge new chunks into it, and save back.
|
||
|
||
Never replaces - only grows (or prunes deleted-file nodes via prune_sources).
|
||
Safe to call repeatedly: existing nodes and edges are preserved.
|
||
root: if given, absolute source_file paths in new_chunks are made relative (#932).
|
||
"""
|
||
graph_path = Path(graph_path)
|
||
if graph_path.exists():
|
||
# Read JSON directly instead of going through node_link_graph().
|
||
# The latter rebuilds an undirected nx.Graph and then enumerating
|
||
# edges() yields endpoints based on node insertion order, which
|
||
# silently flips directional edges (e.g. `calls`) when the callee
|
||
# was inserted before the caller. The _src/_tgt direction-preserving
|
||
# attrs are popped before saving in export.py, so going through the
|
||
# NetworkX round-trip loses direction permanently (#760).
|
||
from graphify.security import check_graph_file_size_cap
|
||
check_graph_file_size_cap(graph_path)
|
||
data = json.loads(graph_path.read_text(encoding="utf-8"))
|
||
links_key = "links" if "links" in data else "edges"
|
||
existing_nodes = list(data.get("nodes", []))
|
||
existing_edges = list(data.get(links_key, []))
|
||
base = [{"nodes": existing_nodes, "edges": existing_edges}]
|
||
else:
|
||
existing_nodes = []
|
||
base = []
|
||
|
||
all_chunks = base + list(new_chunks)
|
||
G = build(all_chunks, directed=directed, dedup=dedup, dedup_llm_backend=dedup_llm_backend, root=root)
|
||
|
||
# Prune nodes and edges from deleted source files
|
||
if prune_sources:
|
||
# Build a set containing both the raw form (matches nodes that kept
|
||
# absolute source_file) and the normalised relative form (matches nodes
|
||
# that were relativised by _norm_source_file at build time).
|
||
# .resolve() handles symlinked roots and redundant ".." / "./" segments
|
||
# so Path.relative_to() succeeds even when the scan root is a symlink.
|
||
# (#1007: manifest absolute paths vs graph relative source_file mismatch)
|
||
_root_str = str(Path(root).resolve()) if root is not None else None
|
||
prune_set: set[str] = set()
|
||
for p in prune_sources:
|
||
if not p:
|
||
continue
|
||
prune_set.add(p)
|
||
norm = _norm_source_file(p, _root_str)
|
||
if norm:
|
||
prune_set.add(norm)
|
||
to_remove = [
|
||
n for n, d in G.nodes(data=True)
|
||
if d.get("source_file") in prune_set
|
||
]
|
||
G.remove_nodes_from(to_remove)
|
||
n_files = len(prune_sources)
|
||
n_nodes = len(to_remove)
|
||
if n_nodes:
|
||
print(
|
||
f"[graphify] Pruned {n_nodes} node(s) from {n_files} deleted source file(s).",
|
||
file=sys.stderr,
|
||
)
|
||
|
||
edges_to_remove = [
|
||
(u, v) for u, v, d in G.edges(data=True)
|
||
if d.get("source_file") in prune_set
|
||
]
|
||
if edges_to_remove:
|
||
G.remove_edges_from(edges_to_remove)
|
||
print(
|
||
f"[graphify] Pruned {len(edges_to_remove)} edge(s) from deleted source file(s).",
|
||
file=sys.stderr,
|
||
)
|
||
|
||
if not n_nodes and not edges_to_remove:
|
||
print(
|
||
f"[graphify] {n_files} source file(s) deleted since last run — "
|
||
f"no matching nodes or edges in graph, already clean.",
|
||
file=sys.stderr,
|
||
)
|
||
|
||
# Safety check: refuse to shrink the graph silently (#479)
|
||
# Skip when dedup or prune_sources is active — shrinkage is intentional there.
|
||
if graph_path.exists() and not dedup and not prune_sources:
|
||
existing_n = len(existing_nodes)
|
||
new_n = G.number_of_nodes()
|
||
if new_n < existing_n:
|
||
raise ValueError(
|
||
f"graphify: build_merge would shrink graph from {existing_n} → {new_n} nodes. "
|
||
f"Pass prune_sources explicitly if you intend to remove nodes."
|
||
)
|
||
|
||
return G
|
||
|
||
|
||
def prefix_graph_for_global(G: nx.Graph, repo_tag: str) -> nx.Graph:
|
||
"""Return a copy of G with all node IDs prefixed with repo_tag::.
|
||
|
||
Labels are preserved unchanged (for display). A 'local_id' attribute
|
||
is added to each node so the original ID can be recovered. Edges are
|
||
rewritten to match the new prefixed IDs. The 'repo' attribute is set
|
||
on every node.
|
||
"""
|
||
relabel = {n: f"{repo_tag}::{n}" for n in G.nodes}
|
||
H = nx.relabel_nodes(G, relabel, copy=True)
|
||
for node, data in H.nodes(data=True):
|
||
data["repo"] = repo_tag
|
||
data.setdefault("local_id", node.split("::", 1)[1])
|
||
return H
|
||
|
||
|
||
def prune_repo_from_graph(G: nx.Graph, repo_tag: str) -> int:
|
||
"""Remove all nodes tagged with repo_tag from G in-place. Returns count removed."""
|
||
to_remove = [n for n, d in G.nodes(data=True) if d.get("repo") == repo_tag]
|
||
G.remove_nodes_from(to_remove)
|
||
return len(to_remove)
|