Blast radius visualization of the `export_asset` function. Click and drag to pan, scroll to zoom.
Overview
To demonstrate the visualization capabilities of the codegen we will generate three different visualizations of PostHog’s open source repository.
Call Trace Visualization
Visualizing the call trace of a function is a great way to understand the flow of a function and for debugging. In this tutorial we will create a call trace visualization of the patch method of the SharingConfigurationViewSet class. View the source code here.
Basic Setup
First, we’ll set up our codebase, graph and configure some basic parameters:
import networkx as nx
from codegen import Codebase
# Initialize codebase
codebase = Codebase("path/to/posthog/")
# Create a directed graph for representing call relationships
G = nx.DiGraph()
# Configuration flags
IGNORE_EXTERNAL_MODULE_CALLS = True # Skip calls to external modules
IGNORE_CLASS_CALLS = False # Include class definition calls
MAX_DEPTH = 10
COLOR_PALETTE = {
"StartFunction": "#9cdcfe", # Light blue - Start Function
"PyFunction": "#a277ff", # Soft purple/periwinkle - PyFunction
"PyClass": "#ffca85", # Warm peach/orange - PyClass
"ExternalModule": "#f694ff" # Bright magenta/pink - ExternalModule
}
Building the Visualization
We’ll create a function that will recursively traverse the call trace of a function and add nodes and edges to the graph:
def create_downstream_call_trace(src_func: Function, depth: int = 0):
"""Creates call graph by recursively traversing function calls
Args:
src_func (Function): Starting function for call graph
depth (int): Current recursion depth
"""
# Prevent infinite recursion
if MAX_DEPTH <= depth:
return
# External modules are not functions
if isinstance(src_func, ExternalModule):
return
# Process each function call
for call in src_func.function_calls:
# Skip self-recursive calls
if call.name == src_func.name:
continue
# Get called function definition
func = call.function_definition
if not func:
continue
# Apply configured filters
if isinstance(func, ExternalModule) and IGNORE_EXTERNAL_MODULE_CALLS:
continue
if isinstance(func, Class) and IGNORE_CLASS_CALLS:
continue
# Generate display name (include class for methods)
if isinstance(func, Class) or isinstance(func, ExternalModule):
func_name = func.name
elif isinstance(func, Function):
func_name = f"{func.parent_class.name}.{func.name}" if func.is_method else func.name
# Add node and edge with metadata
G.add_node(func, name=func_name,
color=COLOR_PALETTE.get(func.__class__.__name__))
G.add_edge(src_func, func, **generate_edge_meta(call))
# Recurse for regular functions
if isinstance(func, Function):
create_downstream_call_trace(func, depth + 1)
We can enrich our edges with metadata about the function calls:
def generate_edge_meta(call: FunctionCall) -> dict:
"""Generate metadata for call graph edges
Args:
call (FunctionCall): Function call information
Returns:
dict: Edge metadata including name and location
"""
return {
"name": call.name,
"file_path": call.filepath,
"start_point": call.start_point,
"end_point": call.end_point,
"symbol_name": "FunctionCall"
}
Visualizing the Graph
Finally, we can visualize our call graph starting from a specific function:
# Get target function to analyze
target_class = codebase.get_class('SharingConfigurationViewSet')
target_method = target_class.get_method('patch')
# Add root node
G.add_node(target_method,
name=f"{target_class.name}.{target_method.name}",
color=COLOR_PALETTE["StartFunction"])
# Build the call graph
create_downstream_call_trace(target_method)
# Render the visualization
codebase.visualize(G)
Take a look
Common Use Cases
The call graph visualization is particularly useful for:
- Understanding complex codebases
- Planning refactoring efforts
- Identifying tightly coupled components
- Analyzing critical paths
- Documenting system architecture
Function Dependency Graph
Understanding symbol dependencies is crucial for maintaining and refactoring code. This tutorial will show you how to create visual dependency graphs using Codegen and NetworkX. We will be creating a dependency graph of the get_query_runner function. View the source code here.
Basic Setup
Building the Dependency Graph
The core function for building our dependency graph:
def create_dependencies_visualization(symbol: Symbol, depth: int = 0):
"""Creates visualization of symbol dependencies
Args:
symbol (Symbol): Starting symbol to analyze
depth (int): Current recursion depth
"""
# Prevent excessive recursion
if depth >= MAX_DEPTH:
return
# Process each dependency
for dep in symbol.dependencies:
dep_symbol = None
# Handle different dependency types
if isinstance(dep, Symbol):
# Direct symbol reference
dep_symbol = dep
elif isinstance(dep, Import):
# Import statement - get resolved symbol
dep_symbol = dep.resolved_symbol if dep.resolved_symbol else None
if dep_symbol:
# Add node with appropriate styling
G.add_node(dep_symbol,
color=COLOR_PALETTE.get(dep_symbol.__class__.__name__,
"#f694ff"))
# Add dependency relationship
G.add_edge(symbol, dep_symbol)
# Recurse unless it's a class (avoid complexity)
if not isinstance(dep_symbol, PyClass):
create_dependencies_visualization(dep_symbol, depth + 1)
Visualizing the Graph
Finally, we can visualize our dependency graph starting from a specific symbol:
# Get target symbol
target_func = codebase.get_function("get_query_runner")
# Add root node
G.add_node(target_func, color=COLOR_PALETTE["StartFunction"])
# Generate dependency graph
create_dependencies_visualization(target_func)
# Render visualization
codebase.visualize(G)
Take a look
Blast Radius visualization
Understanding the impact of code changes is crucial for safe refactoring. A blast radius visualization shows how changes to one function might affect other parts of the codebase by tracing usage relationships. In this tutorial we will create a blast radius visualization of the export_asset function. View the source code here.
Basic Setup
Helper Functions
We’ll create some utility functions to help build our visualization:
# List of HTTP methods to highlight
HTTP_METHODS = ["get", "put", "patch", "post", "head", "delete"]
def generate_edge_meta(usage: Usage) -> dict:
"""Generate metadata for graph edges
Args:
usage (Usage): Usage relationship information
Returns:
dict: Edge metadata including name and location
"""
return {
"name": usage.match.source,
"file_path": usage.match.filepath,
"start_point": usage.match.start_point,
"end_point": usage.match.end_point,
"symbol_name": usage.match.__class__.__name__
}
def is_http_method(symbol: PySymbol) -> bool:
"""Check if a symbol is an HTTP endpoint method
Args:
symbol (PySymbol): Symbol to check
Returns:
bool: True if symbol is an HTTP method
"""
if isinstance(symbol, PyFunction) and symbol.is_method:
return symbol.name in HTTP_METHODS
return False
Building the Blast Radius Visualization
The main function for creating our blast radius visualization:
def create_blast_radius_visualization(symbol: PySymbol, depth: int = 0):
"""Create visualization of symbol usage relationships
Args:
symbol (PySymbol): Starting symbol to analyze
depth (int): Current recursion depth
"""
# Prevent excessive recursion
if depth >= MAX_DEPTH:
return
# Process each usage of the symbol
for usage in symbol.usages:
usage_symbol = usage.usage_symbol
# Determine node color based on type
if is_http_method(usage_symbol):
color = COLOR_PALETTE.get("HTTP_METHOD")
else:
color = COLOR_PALETTE.get(usage_symbol.__class__.__name__, "#f694ff")
# Add node and edge to graph
G.add_node(usage_symbol, color=color)
G.add_edge(symbol, usage_symbol, **generate_edge_meta(usage))
# Recursively process usage symbol
create_blast_radius_visualization(usage_symbol, depth + 1)
Visualizing the Graph
Finally, we can create our blast radius visualization:
# Get target function to analyze
target_func = codebase.get_function('export_asset')
# Add root node
G.add_node(target_func, color=COLOR_PALETTE.get("StartFunction"))
# Build the visualization
create_blast_radius_visualization(target_func)
# Render graph to show impact flow
# Note: a -> b means changes to a will impact b
codebase.visualize(G)
Take a look
What’s Next?
