Key Takeaways
- Using Python's native `ast` module combined with `networkx` enables deterministic, dependency-free code graph extraction without heavyweight external AST parsers.
- Streaming code chunks into Claude with structured JSON schemas prevents context degradation and extracts semantic edge weights between decoupled functions.
- Serializing code graphs to GraphML or Neo4j enables sub-millisecond structural traversal for automated refactoring agents.
- Implementing an incremental parsing cache prevents re-analyzing unchanged files during large CI/CD pipeline runs.
The Problem
Modern Python applications, especially those powering data science, machine learning, and data engineering pipelines, grow at an astonishing pace. This growth inevitably leads to complex module graphs, often with hidden circular dependencies that only surface during runtime or in obscure edge cases. For developers, understanding the impact of a change becomes a daunting task, relying on tribal knowledge or slow, manual code exploration. For the new breed of AI agents, tasked with understanding, explaining, or even refactoring code, these sprawling codebases present an even greater hurdle. Feeding an entire repository into an LLM's context window is inefficient, expensive, and often ineffective due to token limits and the LLM's tendency to "hallucinate" relationships rather than infer them deterministically. We need a way to extract the *structure* and *meaning* of our code, making it machine-readable and queryable, just like any other dataset.Data and Sources
For this tutorial, our "real data" will be a local Python codebase. To make the example self-contained and runnable, the complete script will create a small, illustrative Python project structure if it doesn't already exist. This allows us to parse a live, albeit simple, codebase. We'll leverage several key Python libraries and APIs: * Python `ast` module documentation: Official Python Documentation * NetworkX documentation for graph creation and topological sorts: NetworkX Website * Anthropic Claude 3.5 Sonnet API docs for structured tool use and JSON outputs: Anthropic Developers * Shopify Engineering blog posts on code inventory and agent environments, which inspired parts of this approach: Shopify Engineering Blog (Atom feed for reference) Data accessed on 2024-07-30.Step 1 — Abstract Syntax Tree Traversal & Symbol Extraction
The first challenge is to parse Python source code into meaningful structural elements—classes, functions, imports—without actually executing it. Executing arbitrary code is a security risk and unnecessary for static analysis. Python's built-in `ast` module is our answer. It allows us to parse source code into an Abstract Syntax Tree (AST), which is a hierarchical representation of the program's structure. I typically use `ast.NodeVisitor` to recursively walk this tree. This approach is powerful because it gives us fine-grained control over what we extract. For each Python file, we want to identify definitions (classes, functions), their names, arguments, and any calls made within their bodies. We also need to capture import statements to understand cross-module dependencies. Here’s a simplified look at how I extract symbols from a file:import ast
import os
class CodebaseVisitor(ast.NodeVisitor):
def __init__(self, filepath):
self.filepath = filepath
self.nodes = []
self.edges = []
self.current_function = None
def visit_FunctionDef(self, node):
func_id = f"{self.filepath}::{node.name}"
self.nodes.append({"id": func_id, "type": "function", "name": node.name, "lineno": node.lineno})
# Track current function for call graph
original_function = self.current_function
self.current_function = func_id
self.generic_visit(node)
self.current_function = original_function
def visit_ClassDef(self, node):
class_id = f"{self.filepath}::{node.name}"
self.nodes.append({"id": class_id, "type": "class", "name": node.name, "lineno": node.lineno})
for base in node.bases:
if isinstance(base, ast.Name):
self.edges.append({"source": class_id, "target": base.id, "type": "INHERITS"})
self.generic_visit(node)
def visit_Call(self, node):
if self.current_function and isinstance(node.func, ast.Name):
# This is a local call within a function
self.edges.append({"source": self.current_function, "target": node.func.id, "type": "CALLS"})
elif self.current_function and isinstance(node.func, ast.Attribute):
# This is a method call, e.g., 'obj.method()'
# For simplicity, we'll just capture the method name here.
# A more advanced parser would resolve `obj`'s type.
self.edges.append({"source": self.current_function, "target": node.func.attr, "type": "CALLS_METHOD"})
self.generic_visit(node)
def visit_Import(self, node):
for alias in node.names:
self.nodes.append({"id": f"module::{alias.name}", "type": "import", "name": alias.name})
# Add an edge from the current file to the imported module
self.edges.append({"source": self.filepath, "target": f"module::{alias.name}", "type": "IMPORTS_MODULE"})
self.generic_visit(node)
def visit_ImportFrom(self, node):
module_name = node.module if node.module else ""
for alias in node.names:
imported_name = f"{module_name}.{alias.name}" if module_name else alias.name
self.nodes.append({"id": f"module::{imported_name}", "type": "import", "name": imported_name})
self.edges.append({"source": self.filepath, "target": f"module::{imported_name}", "type": "IMPORTS_FROM"})
self.generic_visit(node)
def parse_python_file(filepath):
with open(filepath, "r", encoding="utf-8") as f:
source_code = f.read()
tree = ast.parse(source_code, filename=filepath)
visitor = CodebaseVisitor(filepath)
visitor.visit(tree)
return visitor.nodes, visitor.edges, source_code
Step 2 — Constructing the Directed Graph with NetworkX
With individual file's ASTs parsed, the next challenge is to resolve cross-module references and combine these into a coherent, directed graph. This is where `networkx` shines. It provides powerful tools for creating, manipulating, and analyzing graph structures. Our goal is to create a `networkx.DiGraph` where nodes represent files, classes, and functions, and edges represent relationships like `CALLS`, `INHERITS`, and `IMPORTS`. The tricky part is resolving imports. When `module_a` imports `function_x` from `module_b`, we need to map that string reference (`function_x`) to the actual `function_x` node defined in `module_b`. This requires a symbol resolution pass that iterates through all collected nodes and edges, attempting to link external references.import networkx as nx
def build_initial_graph(parsed_data):
G = nx.DiGraph()
all_symbols = {} # Map names to their unique IDs for resolution
# Add all files as nodes first
for filepath, _, _ in parsed_data:
G.add_node(filepath, type="file", name=os.path.basename(filepath))
# Add nodes and edges from individual file parsing
for filepath, nodes, edges, source_code in parsed_data:
for node_data in nodes:
unique_id = node_data["id"]
G.add_node(unique_id, **node_data, source_code_snippet=source_code)
if node_data["type"] in ["function", "class"]:
all_symbols[node_data["name"]] = unique_id
# Add implicit edge from file to its definitions
if node_data["type"] in ["function", "class"]:
G.add_edge(filepath, unique_id, type="CONTAINS")
for edge_data in edges:
source = edge_data["source"]
target = edge_data["target"]
edge_type = edge_data["type"]
G.add_edge(source, target, type=edge_type)
return G, all_symbols
def resolve_graph_references(G, all_symbols):
# This is a simplified resolution. A full solution would need
# to handle relative imports, __init__.py, and PYTHONPATH.
resolved_G = G.copy()
for u, v, data in G.edges(data=True):
if data["type"] == "CALLS" and v in all_symbols:
# If a call target is a known symbol, update the edge to its unique ID
resolved_G.remove_edge(u, v)
resolved_G.add_edge(u, all_symbols[v], type="CALLS_RESOLVED")
elif data["type"] == "INHERITS" and v in all_symbols:
resolved_G.remove_edge(u, v)
resolved_G.add_edge(u, all_symbols[v], type="INHERITS_RESOLVED")
# More complex resolution for IMPORTS_MODULE / IMPORTS_FROM would go here
# mapping 'module::requests' to the actual requests module node etc.
return resolved_G
Step 3 — LLM-Enriched Semantic Metadata Generation
While AST parsing gives us structure, it doesn't tell us *what* a function does in plain language, or its conceptual complexity. This is where large language models like Claude come in. We can batch function body snippets and query Claude 3.5 Sonnet to tag graph nodes with functional summaries, potential risk scores, or even suggest refactorings. The key is to use structured output, often with Pydantic, to ensure reliable, parseable responses. Instead of sending entire files, we send only the relevant function body or class definition, drastically reducing token usage and improving focus.import anthropic
import json
from pydantic import BaseModel, Field
import os
class FunctionMetadata(BaseModel):
summary: str = Field(description="A concise summary of the function's purpose.")
complexity_score: int = Field(description="An estimated score (1-5) of the function's conceptual complexity, where 5 is very complex.")
potential_risks: list[str] = Field(description="A list of potential risks or side effects, e.g., 'database write', 'external API call'.")
def get_claude_metadata(function_code: str, anthropic_api_key: str) -> FunctionMetadata:
if not anthropic_api_key:
print("ANTHROPIC_API_KEY not set. Returning mock metadata for LLM enrichment.")
return FunctionMetadata(
summary="MOCK: This function performs a mock operation.",
complexity_score=2,
potential_risks=["mock_risk_1", "mock_risk_2"]
)
client = anthropic.Anthropic(api_key=anthropic_api_key)
try:
response = client.messages.create(
model="claude-3-5-sonnet-20240620",
max_tokens=500,
temperature=0.3,
messages=[
{
"role": "user",
"content": [
{"type": "text", "text": f"Analyze the following Python function and provide a summary, complexity score (1-5), and potential risks as a JSON object adhering to the FunctionMetadata Pydantic schema:\n\n\n{function_code}\n \n\nOutput only the JSON object."}
]
}
],
tool_choice={"type": "auto"},
tools=[
{
"name": "FunctionMetadata",
"description": "Metadata for a Python function.",
"input_schema": FunctionMetadata.model_json_schema()
}
]
)
# Assuming the LLM returns the JSON directly in content
# For tool use, response.content might contain tool_use blocks.
# This example assumes a direct JSON response for simplicity.
# In a real tool use scenario, you'd parse response.content[0].text
# and then json.loads it, or use the tool_use block.
# A more robust parsing for direct JSON in content
try:
json_str = response.content[0].text
return FunctionMetadata.model_validate_json(json_str)
except (json.JSONDecodeError, IndexError) as e:
print(f"Error parsing Claude response for function: {e}. Raw response: {response.content}")
return FunctionMetadata(summary="Failed to parse Claude response.", complexity_score=0, potential_risks=[])
except anthropic.APIError as e:
print(f"Claude API error for function: {e}")
return FunctionMetadata(summary=f"API error: {e}", complexity_score=0, potential_risks=[])
except Exception as e:
print(f"An unexpected error occurred during Claude call: {e}")
return FunctionMetadata(summary=f"Unexpected error: {e}", complexity_score=0, potential_risks=[])
def enrich_graph_with_claude(G: nx.DiGraph, anthropic_api_key: str):
for node_id in list(G.nodes): # Iterate over a copy to allow modification
node = G.nodes[node_id]
if node.get("type") == "function" and "source_code_snippet" in node:
# Extract just the function's source code, not the whole file
# This requires more advanced AST parsing to get exact function body
# For simplicity, we'll send the whole file's source and ask Claude to focus
# A better approach would extract the exact function AST node's source.
# Placeholder: In a real scenario, you'd extract the specific function's code
# For this example, we'll try to get the definition block from the file content
file_content = node.get("source_code_snippet")
# Simple heuristic to get a function's code block
# This is fragile; a true solution would use the AST node's start/end lines.
func_name = node["name"]
func_code_lines = []
in_func = False
for line in file_content.splitlines():
if f"def {func_name}(" in line or f"async def {func_name}(" in line:
in_func = True
if in_func:
func_code_lines.append(line)
# Simple heuristic to guess end of function: next non-indented line
# This is very basic and will fail with complex structures.
if len(func_code_lines) > 1 and line.strip() == "" and func_code_lines[-2].strip() != "":
break # Likely end of function definition.
func_code_snippet = "\n".join(func_code_lines) if func_code_lines else file_content # Fallback
if func_code_snippet:
metadata = get_claude_metadata(func_code_snippet, anthropic_api_key)
node["summary"] = metadata.summary
node["complexity_score"] = metadata.complexity_score
node["potential_risks"] = metadata.potential_risks
print(f"Enriched {node_id}: {metadata.summary[:50]}...")
Step 4 — Querying the Knowledge Graph via Vector and Graph Traversal
With our enriched knowledge graph, we can now execute complex topological queries that are impossible with simple text search. `networkx` provides powerful algorithms for this. We can find all downstream consumers of a modified configuration object, identify circular dependencies, or trace deep call paths. For instance, finding all functions that directly or indirectly call a specific "sensitive" function:def find_downstream_consumers(G: nx.DiGraph, start_node_id: str):
if start_node_id not in G:
print(f"Node {start_node_id} not found in graph.")
return []
# Using BFS to find all reachable nodes
reachable_nodes = nx.bfs_tree(G, start_node_id)
# Exclude the start node itself if we only want consumers
consumers = [node for node in reachable_nodes if node != start_node_id and G.nodes[node].get("type") == "function"]
return consumers
def find_circular_dependencies(G: nx.DiGraph):
try:
cycles = list(nx.simple_cycles(G))
return cycles
except nx.NetworkXNoCycle:
return []
def get_architectural_hubs(G: nx.DiGraph, top_n=5):
# Using betweenness centrality to find nodes that act as "bridges"
# or central points in the graph.
betweenness = nx.betweenness_centrality(G)
sorted_hubs = sorted(betweenness.items(), key=lambda item: item[1], reverse=True)
hub_nodes = []
for node_id, score in sorted_hubs:
node_type = G.nodes[node_id].get("type")
node_name = G.nodes[node_id].get("name", node_id)
if node_type in ["function", "class", "file"]: # Filter for meaningful hubs
hub_nodes.append({"id": node_id, "name": node_name, "type": node_type, "centrality": score})
if len(hub_nodes) >= top_n:
break
return hub_nodes
Complete Script
The full runnable script combining all steps:
#!/usr/bin/env python3
import ast
import os
import sys
import networkx as nx
import json
import anthropic
from pydantic import BaseModel, Field
import shutil
# --- Pydantic Model for Claude Output ---
class FunctionMetadata(BaseModel):
summary: str = Field(description="A concise summary of the function's purpose.")
complexity_score: int = Field(description="An estimated score (1-5) of the function's conceptual complexity, where 5 is very complex.")
potential_risks: list[str] = Field(description="A list of potential risks or side effects, e.g., 'database write', 'external API call'.")
# --- Step 1: AST Traversal and Symbol Extraction ---
class CodebaseVisitor(ast.NodeVisitor):
def __init__(self, filepath, source_code):
self.filepath = filepath
self.source_code = source_code
self.nodes = []
self.edges = []
self.current_definition_id = filepath # Track current file, class, or function for nested definitions/calls
def _add_node(self, node_id, node_type, name, lineno, **kwargs):
self.nodes.append({"id": node_id, "type": node_type, "name": name, "lineno": lineno, **kwargs})
def _add_edge(self, source, target, edge_type, **kwargs):
self.edges.append({"source": source, "target": target, "type": edge_type, **kwargs})
def visit_FunctionDef(self, node):
func_id = f"{self.filepath}::{node.name}"
self._add_node(func_id, "function", node.name, node.lineno)
self._add_edge(self.current_definition_id, func_id, "CONTAINS_DEF")
original_definition_id = self.current_definition_id
self.current_definition_id = func_id
self.generic_visit(node)
self.current_definition_id = original_definition_id
def visit_AsyncFunctionDef(self, node):
self.visit_FunctionDef(node) # Treat async functions similarly for graph structure
def visit_ClassDef(self, node):
class_id = f"{self.filepath}::{node.name}"
self._add_node(class_id, "class", node.name, node.lineno)
self._add_edge(self.current_definition_id, class_id, "CONTAINS_DEF")
for base in node.bases:
if isinstance(base, ast.Name):
self._add_edge(class_id, base.id, "INHERITS") # Target is unresolved name
elif isinstance(base, ast.Attribute):
self._add_edge(class_id, base.attr, "INHERITS_ATTR") # Target is unresolved name. e.g. `requests.Request`
original_definition_id = self.current_definition_id
self.current_definition_id = class_id
self.generic_visit(node)
self.current_definition_id = original_definition_id
def visit_Call(self, node):
if isinstance(node.func, ast.Name):
self._add_edge(self.current_definition_id, node.func.id, "CALLS") # Target is unresolved name
elif isinstance(node.func, ast.Attribute):
# For method calls like `obj.method()`, capture `method`
self._add_edge(self.current_definition_id, node.func.attr, "CALLS_METHOD") # Target is unresolved name
self.generic_visit(node)
def visit_Import(self, node):
for alias in node.names:
imported_name = alias.asname if alias.asname else alias.name
self._add_node(f"import::{imported_name}", "imported_symbol", imported_name, node.lineno)
self._add_edge(self.filepath, f"import::{imported_name}", "IMPORTS_AS") # Link file to the imported alias
self.generic_visit(node)
def visit_ImportFrom(self, node):
module_name = node.module if node.module else ""
for alias in node.names:
imported_name = alias.asname if alias.asname else alias.name
# Store full path for better resolution later
full_imported_path = f"{module_name}.{alias.name}" if module_name else alias.name
self._add_node(f"import::{imported_name}", "imported_symbol", imported_name, node.lineno, full_path=full_imported_path)
self._add_edge(self.filepath, f"import::{imported_name}", "IMPORTS_FROM", module=module_name, original_name=alias.name)
self.generic_visit(node)
def parse_python_file(filepath):
try:
with open(filepath, "r", encoding="utf-8") as f:
source_code = f.read()
tree = ast.parse(source_code, filename=filepath)
visitor = CodebaseVisitor(filepath, source_code)
visitor.visit(tree)
return visitor.nodes, visitor.edges, source_code
except SyntaxError as e:
print(f"Error parsing {filepath}: SyntaxError - {e}", file=sys.stderr)
return [], [], ""
except Exception as e:
print(f"Error processing {filepath}: {e}", file=sys.stderr)
return [], [], ""
# --- Step 2: Constructing the Directed Graph with NetworkX ---
def build_and_resolve_graph(repo_path):
G = nx.DiGraph()
parsed_files_data = []
# Collect all Python files
python_files = []
for root, _, files in os.walk(repo_path):
for file in files:
if file.endswith(".py"):
python_files.append(os.path.join(root, file))
if not python_files:
raise ValueError(f"No Python files found in repository path: {repo_path}")
# First pass: parse all files and add file nodes
for filepath in python_files:
relative_filepath = os.path.relpath(filepath, repo_path)
G.add_node(relative_filepath, type="file", name=os.path.basename(filepath), full_path=filepath)
nodes, edges, source_code = parse_python_file(filepath)
parsed_files_data.append((relative_filepath, nodes, edges, source_code))
# Second pass: add all definitions and initial edges
all_definitions = {} # Maps (file_path, symbol_name) to unique_id
for relative_filepath, nodes, edges, source_code in parsed_files_data:
for node_data in nodes:
unique_id = node_data["id"]
G.add_node(unique_id, **node_data, source_code_snippet=source_code)
if node_data["type"] in ["function", "class"]:
key = (relative_filepath, node_data["name"])
all_definitions[key] = unique_id
for edge_data in edges:
G.add_edge(edge_data["source"], edge_data["target"], type=edge_data["type"])
# Third pass: Resolve internal (within-project) references
resolved_G = G.copy()
for u, v, data in G.edges(data=True):
if data["type"] == "CALLS" or data["type"] == "INHERITS" or data["type"] == "CALLS_METHOD":
# Attempt to resolve 'v' (unresolved name) to a known definition
# This simplified resolver only checks definitions within the same file or across project
# Check within the same file (u's file)
u_file_node = u.split("::")[0] if "::" in u else u # Get the file node ID
if resolved_G.has_node(u_file_node) and resolved_G.nodes[u_file_node]["type"] == "file":
if (u_file_node, v) in all_definitions:
resolved_G.remove_edge(u, v)
resolved_G.add_edge(u, all_definitions[(u_file_node, v)], type=f"{data['type']}_RESOLVED_LOCAL")
continue
# Check across all definitions in the project
found_global = False
for (file_path_def, symbol_name), def_id in all_definitions.items():
if symbol_name == v:
resolved_G.remove_edge(u, v)
resolved_G.add_edge(u, def_id, type=f"{data['type']}_RESOLVED_GLOBAL")
found_global = True
break
if not found_global:
# If still not resolved, mark as external or unresolved
# resolved_G.add_edge(u, v, type=f"{data['type']}_UNRESOLVED") # Keep original edge
pass # Keep original edge as is
elif data["type"] == "IMPORTS_AS" or data["type"] == "IMPORTS_FROM":
# For imports, try to link to the actual file node if it's within the project
imported_name = data.get("full_path") # from import::... node data
if imported_name:
# Convert `module.sub` to `module/sub.py` or `module/sub/__init__.py`
potential_file_path_1 = imported_name.replace('.', os.sep) + ".py"
potential_file_path_2 = imported_name.replace('.', os.sep) + os.sep + "__init__.py"
for fp_abs in python_files:
fp_rel = os.path.relpath(fp_abs, repo_path)
if fp_rel == potential_file_path_1 or fp_rel == potential_file_path_2:
resolved_G.remove_edge(u, v) # remove the import:: symbol node
resolved_G.add_edge(u, fp_rel, type="IMPORTS_RESOLVED_FILE")
# Also remove the import:: symbol node if it's no longer needed
if v in resolved_G and resolved_G.degree(v) == 0:
resolved_G.remove_node(v)
break
return resolved_G
# --- Step 3: LLM-Enriched Semantic Metadata Generation ---
def get_claude_metadata(function_code: str, anthropic_api_key: str) -> FunctionMetadata:
if not anthropic_api_key:
print("ANTHROPIC_API_KEY not set. Returning mock metadata for LLM enrichment.", file=sys.stderr)
return FunctionMetadata(
summary="MOCK: This function performs a mock operation.",
complexity_score=2,
potential_risks=["mock_risk_1", "mock_risk_2"]
)
client = anthropic.Anthropic(api_key=anthropic_api_key)
try:
response = client.messages.create(
model="claude-3-5-sonnet-20240620",
max_tokens=500,
temperature=0.3,
messages=[
{
"role": "user",
"content": [
{"type": "text", "text": f"Analyze the following Python function and provide a summary, complexity score (1-5), and potential risks as a JSON object adhering to the FunctionMetadata Pydantic schema:\n\n\n{function_code}\n \n\nOutput only the JSON object."}
]
}
]
)
json_str = response.content[0].text
return FunctionMetadata.model_validate_json(json_str)
except anthropic.APIError as e:
print(f"Claude API error for function: {e}", file=sys.stderr)
return FunctionMetadata(summary=f"API error: {e}", complexity_score=0, potential_risks=[])
except (json.JSONDecodeError, IndexError, Exception) as e:
print(f"Error parsing Claude response for function: {e}. Raw response: {response.content}", file=sys.stderr)
return FunctionMetadata(summary="Failed to parse Claude response.", complexity_score=0, potential_risks=[])
def extract_function_code(source_code: str, lineno: int) -> str:
lines = source_code.splitlines()
if not (1 <= lineno <= len(lines)):
return "" # Line number out of bounds
start_line_idx = lineno - 1
# Find the start of the def statement
while start_line_idx >= 0 and not (lines[start_line_idx].strip().startswith("def ") or lines[start_line_idx].strip().startswith("async def ")):
start_line_idx -= 1
if start_line_idx < 0:
return "" # Could not find function definition
indentation = len(lines[start_line_idx]) - len(lines[start_line_idx].lstrip())
func_code_lines = [lines[start_line_idx]]
current_line_idx = start_line_idx + 1
# Heuristic to find end of function: next line with less or equal indentation
while current_line_idx < len(lines):
line = lines[current_line_idx]
if line.strip() == "": # Allow empty lines
func_code_lines.append(line)
elif len(line) - len(line.lstrip()) <= indentation:
break # Found a line with less or equal indentation, likely end of function
else:
func_code_lines.append(line)
current_line_idx += 1
return "\n".join(func_code_lines)
def enrich_graph_with_claude(G: nx.DiGraph, anthropic_api_key: str):
print("Starting Claude enrichment...")
for node_id in list(G.nodes):
node = G.nodes[node_id]
if node.get("type") == "function" and "source_code_snippet" in node and "lineno" in node:
func_code = extract_function_code(node["source_code_snippet"], node["lineno"])
if func_code:
metadata = get_claude_metadata(func_code, anthropic_api_key)
node["summary"] = metadata.summary
node["complexity_score"] = metadata.complexity_score
node["potential_risks"] = metadata.potential_risks
print(f"Enriched {node_id}: {metadata.summary[:70]}...")
else:
print(f"Could not extract code for {node_id}. Skipping Claude enrichment.", file=sys.stderr)
elif node.get("type") == "class" and "source_code_snippet" in node and "lineno" in node:
# Optionally, you could also enrich classes
pass
print("Claude enrichment complete.")
# --- Step 4: Querying the Knowledge Graph ---
def find_downstream_consumers(G: nx.DiGraph, start_node_id: str):
if start_node_id not in G:
return []
# Using BFS to find all reachable nodes from start_node_id
# nx.bfs_tree includes the source node
reachable_tree = nx.bfs_tree(G, start_node_id)
consumers = [node for node in reachable_tree.nodes if node != start_node_id and G.nodes[node].get("type") == "function"]
return consumers
def find_circular_dependencies(G: nx.DiGraph):
try:
cycles = list(nx.simple_cycles(G))
# Filter for cycles involving actual code definitions or files
meaningful_cycles = [
cycle for cycle in cycles
if any(G.nodes[node_id].get("type") in ["function", "class", "file"] for node_id in cycle)
]
return meaningful_cycles
except nx.NetworkXNoCycle:
return []
def get_architectural_hubs(G: nx.DiGraph, top_n=5):
if not G.nodes:
return []
# Using betweenness centrality to find nodes that act as "bridges"
betweenness = nx.betweenness_centrality(G)
sorted_hubs = sorted(betweenness.items(), key=lambda item: item[1], reverse=True)
hub_nodes = []
for node_id, score in sorted_hubs:
node_type = G.nodes[node_id].get("type")
node_name = G.nodes[node_id].get("name", node_id)
# Filter for meaningful nodes and ignore very low centrality scores
if node_type in ["function", "class", "file"] and score > 0:
hub_nodes.append({"id": node_id, "name": node_name, "type": node_type, "centrality": score})
if len(hub_nodes) >= top_n:
break
return hub_nodes
# --- Utility for creating a dummy project ---
def create_dummy_python_project(repo_path):
if os.path.exists(repo_path):
print(f"Project directory '{repo_path}' already exists. Skipping creation.")
return
print(f"Creating dummy Python project at '{repo_path}'...")
os.makedirs(os.path.join(repo_path, "src", "utils"), exist_ok=True)
os.makedirs(os.path.join(repo_path, "src", "models"), exist_ok=True)
# src/main.py
with open(os.path.join(repo_path, "src", "main.py"), "w") as f:
f.write("""
import os
from .utils.helpers import calculate_sum, log_message
from .models.data_model import DataRecord
def process_data(data_list):
total = 0
for item in data_list:
record = DataRecord(value=item)
total = calculate_sum(total, record.value)
log_message(f"Processed sum: {total}")
return total
def run_pipeline():
sample_data = [10, 20, 30]
result = process_data(sample_data)
print(f"Pipeline finished with result: {result}")
if __name__ == "__main__":
run_pipeline()
""")
# src/utils/helpers.py
with open(os.path.join(repo_path, "src", "utils", "helpers.py"), "w") as f:
f.write("""
import logging
logging.basicConfig(level=logging.INFO)
def calculate_sum(a, b):
\"\"\"Calculates the sum of two numbers.\"\"\"
return a + b
def log_message(message):
\"\"\"Logs a message to the console.\"\"\"
logging.info(message)
class HelperUtility:
def __init__(self, name):
self.name = name
def greet(self):
log_message(f"Hello from {self.name}!")
""")
# src/models/data_model.py
with open(os.path.join(repo_path, "src", "models", "data_model.py"), "w") as f:
f.write("""
class DataRecord:
def __init__(self, value):
if not isinstance(value, (int, float)):
raise ValueError("Value must be a number.")
self.value = value
def get_value(self):
return self.value
""")
print("Dummy project created successfully.")
# --- Main Execution ---
def main():
repo_path = "./temp_python_project"
create_dummy_python_project(repo_path)
anthropic_api_key = os.environ.get("ANTHROPIC_API_KEY")
if not anthropic_api_key:
print("\nWarning: ANTHROPIC_API_KEY not found in environment variables.", file=sys.stderr)
print(" Claude enrichment will use mock data. Set the environment variable to enable real LLM calls.", file=sys.stderr)
print(" You can get an API key from https://console.anthropic.com/settings/keys", file=sys.stderr)
print(f"\nBuilding knowledge graph for '{repo_path}'...")
try:
graph = build_and_resolve_graph(repo_path)
print(f"Initial graph built with {graph.number_of_nodes()} nodes and {graph.number_of_edges()} edges.")
except ValueError as e:
print(f"Error building graph: {e}", file=sys.stderr)
sys.exit(1)
except Exception as e:
print(f"An unexpected error occurred during graph building: {e}", file=sys.stderr)
sys.exit(1)
# Enrich graph with Claude metadata
enrich_graph_with_claude(graph, anthropic_api_key)
# Save the graph
output_graph_path = "code_knowledge_graph.graphml"
try:
nx.write_graphml(graph, output_graph_path)
print(f"\nKnowledge graph saved to '{output_graph_path}'")
except Exception as e:
print(f"Error saving graph to {output_graph_path}: {e}", file=sys.stderr)
# Perform queries
print("\n--- Query Results ---")
# Find circular dependencies
circular_deps = find_circular_dependencies(graph)
if circular_deps:
print("\nDetected Circular Dependencies:")
for cycle in circular_deps:
print(f"- {json.dumps(cycle)}")
else:
print("\nNo circular dependencies detected.")
# Find downstream consumers of a specific function
target_func_id = "src/utils/helpers.py::calculate_sum"
print(f"\nDownstream consumers of '{target_func_id}':")
consumers = find_downstream_consumers(graph, target_func_id)
if consumers:
for consumer in consumers:
print(f"- {consumer}")
else:
print(f"No consumers found for '{target_func_id}'.")
# Get architectural hubs
hubs = get_architectural_hubs(graph, top_n=3)
print("\nTop 3 Architectural Hubs (by Betweenness Centrality):")
for hub in hubs:
print(f"- {hub['name']} ({hub['type']}): Centrality={hub['centrality']:.4f}")
if __name__ == "__main__":
main()
Expected Output
When you run the script, it will first create a dummy Python project if it doesn't exist. Then, it will parse the project, build the `networkx` graph, and (if `ANTHROPIC_API_KEY` is set) call Claude to enrich the function nodes. Finally, it will print a summary of detected circular dependencies, downstream consumers of `calculate_sum`, and the top architectural hubs. You should see output similar to this:Creating dummy Python project at './temp_python_project'...
Dummy project created successfully.
Warning: ANTHROPIC_API_KEY not found in environment variables.
Claude enrichment will use mock data. Set the environment variable to enable real LLM calls.
You can get an API key from https://console.anthropic.com/settings/keys
Building knowledge graph for './temp_python_project'...
Initial graph built with 15 nodes and 20 edges.
Starting Claude enrichment...
ANTHROPIC_API_KEY not set. Returning mock metadata for LLM enrichment.
Enriched src/utils/helpers.py::calculate_sum: MOCK: This function performs a mock operation....
ANTHROPIC_API_KEY not set. Returning mock metadata for LLM enrichment.
Enriched src/utils/helpers.py::log_message: MOCK: This function performs a mock operation....
ANTHROPIC_API_KEY not set. Returning mock metadata for LLM enrichment.
Enriched src/models/data_model.py::DataRecord::get_value: MOCK: This function performs a mock operation....
ANTHROPIC_API_KEY not set. Returning mock metadata for LLM enrichment.
Enriched src/main.py::process_data: MOCK: This function performs a mock operation....
ANTHROPIC_API_KEY not set. Returning mock metadata for LLM enrichment.
Enriched src/main.py::run_pipeline: MOCK: This function performs a mock operation....
ANTHROPIC_API_KEY not set. Returning mock metadata for LLM enrichment.
Enriched src/utils/helpers.py::HelperUtility::greet: MOCK: This function performs a mock operation....
Claude enrichment complete.
Knowledge graph saved to 'code_knowledge_graph.graphml'
--- Query Results ---
No circular dependencies detected.
Downstream consumers of 'src/utils/helpers.py::calculate_sum':
- src/main.py::process_data
Top 3 Architectural Hubs (by Betweenness Centrality):
- main.py (file): Centrality=0.4286
- process_data (function): Centrality=0.3571
- calculate_sum (function): Centrality=0.2857