Skip to content

Commit bf0f077

Browse files
vishalshenoyvishalshenoy
authored
Visualization Examples (#39)
* . * . * . * . * . * new folder name * . * Automated pre-commit update * fixes
1 parent 6975944 commit bf0f077

File tree

5 files changed

+596
-0
lines changed

5 files changed

+596
-0
lines changed

‎examples/visualize_codebases/README.md‎

Lines changed: 166 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,166 @@
1+
# Codebase Relationship Visualizations
2+
3+
This set of examples demonstrates four different approaches to visualizing code relationships using Codegen. Each visualization script creates a graph to help developers understand different aspects of code structure and dependencies.
4+
5+
## Visualization Types
6+
7+
### 1. Function Call Relationships (`call_trace.py`)
8+
Traces downstream function call relationships from a target method. This visualization is particularly useful for understanding the flow of execution and identifying complex call chains that might need optimization or refactoring.
9+
10+
> [!NOTE]
11+
> View the graph-based visualization created by this script on the `PostHog/posthog` repository [here](https://www.codegen.sh/codemod/6a34b45d-c8ad-422e-95a8-46d4dc3ce2b0/public/diff).
12+
13+
```python
14+
def create_downstream_call_trace(src_func: Function, depth: int = 0):
15+
"""Creates call graph for parent function by recursively traversing all function calls"""
16+
if MAX_DEPTH <= depth:
17+
return
18+
if isinstance(src_func, ExternalModule):
19+
return
20+
21+
for call in src_func.function_calls:
22+
# Skip recursive calls
23+
if call.name == src_func.name:
24+
continue
25+
26+
func = call.function_definition
27+
if not func:
28+
continue
29+
30+
# Add node and edge to graph with metadata
31+
G.add_node(func, name=func_name, color=COLOR_PALETTE.get(func.__class__.__name__))
32+
G.add_edge(src_func, func, **generate_edge_meta(call))
33+
34+
# Recurse for nested calls
35+
if isinstance(func, Function):
36+
create_downstream_call_trace(func, depth + 1)
37+
```
38+
39+
### 2. Symbol Dependencies (`dependency_trace.py`)
40+
Maps symbol dependencies throughout the codebase. This helps developers identify tightly coupled components and understand the impact of modifying shared dependencies, making it easier to plan architectural changes.
41+
42+
> [!NOTE]
43+
> View the graph-based visualization created by this script on the `PostHog/posthog` repository [here](codegen.sh/codemod/f6c63e40-cc20-4b91-a6c7-e5cbd736ce0d/public/diff).
44+
45+
```python
46+
def create_dependencies_visualization(symbol: Symbol, depth: int = 0):
47+
"""Creates a visualization of symbol dependencies in the codebase"""
48+
if depth >= MAX_DEPTH:
49+
return
50+
51+
for dep in symbol.dependencies:
52+
dep_symbol = None
53+
if isinstance(dep, Symbol):
54+
dep_symbol = dep
55+
elif isinstance(dep, Import):
56+
dep_symbol = dep.resolved_symbol if dep.resolved_symbol else None
57+
58+
if dep_symbol:
59+
G.add_node(dep_symbol, color=COLOR_PALETTE.get(dep_symbol.__class__.__name__, "#f694ff"))
60+
G.add_edge(symbol, dep_symbol)
61+
62+
if not isinstance(dep_symbol, Class):
63+
create_dependencies_visualization(dep_symbol, depth + 1)
64+
```
65+
66+
### 3. Function Blast Radius (`blast_radius.py`)
67+
Shows the impact radius of potential changes. This visualization is invaluable for risk assessment before refactoring, as it reveals all the code paths that could be affected by modifying a particular function or symbol.
68+
69+
> [!NOTE]
70+
> View the graph-based visualization created by this script on the `PostHog/posthog` repository [here](codegen.sh/codemod/02f11ebe-6a3a-4687-b31d-2d6bc6a04f3c/public/diff).
71+
72+
```python
73+
def create_blast_radius_visualization(symbol: PySymbol, depth: int = 0):
74+
"""Recursively build a graph visualization showing how a symbol is used"""
75+
if depth >= MAX_DEPTH:
76+
return
77+
78+
for usage in symbol.usages:
79+
usage_symbol = usage.usage_symbol
80+
81+
# Color code HTTP methods specially
82+
if is_http_method(usage_symbol):
83+
color = COLOR_PALETTE.get("HTTP_METHOD")
84+
else:
85+
color = COLOR_PALETTE.get(usage_symbol.__class__.__name__, "#f694ff")
86+
87+
G.add_node(usage_symbol, color=color)
88+
G.add_edge(symbol, usage_symbol, **generate_edge_meta(usage))
89+
90+
create_blast_radius_visualization(usage_symbol, depth + 1)
91+
```
92+
93+
### 4. Class Method Relationships (`method_relationships.py`)
94+
Creates a comprehensive view of class method interactions. This helps developers understand class cohesion, identify potential god classes, and spot opportunities for breaking down complex classes into smaller, more manageable components.
95+
96+
> [!NOTE]
97+
> View the graph-based visualization created by this script on the `modal-labs/modal-client` repository [here](https://www.codegen.sh/codemod/66e2e195-ceec-4935-876a-ed4cfc1731c7/public/diff).
98+
99+
```python
100+
def graph_class_methods(target_class: Class):
101+
"""Creates a graph visualization of all methods in a class and their call relationships"""
102+
G.add_node(target_class, color=COLOR_PALETTE["StartClass"])
103+
104+
# Add all methods as nodes
105+
for method in target_class.methods:
106+
method_name = f"{target_class.name}.{method.name}"
107+
G.add_node(method, name=method_name, color=COLOR_PALETTE["StartMethod"])
108+
visited.add(method)
109+
G.add_edge(target_class, method)
110+
111+
# Create call traces for each method
112+
for method in target_class.methods:
113+
create_downstream_call_trace(method)
114+
```
115+
116+
## Common Features
117+
118+
All visualizations share these characteristics:
119+
120+
1. **Configurable Depth**
121+
- MAX_DEPTH setting controls recursion
122+
- Prevents infinite loops in circular references
123+
124+
2. **Color Coding**
125+
```python
126+
COLOR_PALETTE = {
127+
"StartFunction": "#9cdcfe", # Entry point
128+
"PyFunction": "#a277ff", # Regular functions
129+
"PyClass": "#ffca85", # Classes
130+
"ExternalModule": "#f694ff" # External calls
131+
}
132+
```
133+
134+
3. **Edge Metadata**
135+
- Tracks file paths
136+
- Creates data object for visualization
137+
## Running the Visualizations
138+
139+
```bash
140+
# Install dependencies
141+
pip install codegen networkx
142+
143+
# Run any visualization script
144+
python call_trace.py # Function call relationships
145+
python dependency_trace.py # Symbol dependencies
146+
python blast_radius.py # Function blast radius
147+
python method_relationships.py # Class method relationships
148+
```
149+
150+
Each script will:
151+
1. Initialize the codebase
152+
2. Create the appropriate graph for the relationship
153+
3. Generate visualization data
154+
155+
## View Results
156+
157+
After running a script, you'll get a graph object containing node and edge relationships. You can view an interactive visualization of the graph through the links above pointing to codegen.sh.
158+
159+
## Learn More
160+
161+
- [Codebase Visualization Documentation](https://docs.codegen.com/tutorials/codebase-visualization)
162+
- [Codegen Documentation](https://docs.codegen.com)
163+
164+
## Contributing
165+
166+
Feel free to submit issues and any enhancement requests!

‎examples/visualize_codebases/blast_radius.py‎

Lines changed: 119 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,119 @@
1+
import codegen
2+
from codegen import Codebase
3+
from codegen.sdk.enums import ProgrammingLanguage
4+
import networkx as nx
5+
from codegen.sdk.python.symbol import PySymbol
6+
from codegen.sdk.python.function import PyFunction
7+
from codegen.sdk.core.dataclasses.usage import Usage
8+
9+
# Create a directed graph for visualizing relationships between code elements
10+
G = nx.DiGraph()
11+
12+
# Maximum depth to traverse in the call graph to prevent infinite recursion
13+
MAX_DEPTH = 5
14+
15+
# Define colors for different types of nodes in the visualization
16+
COLOR_PALETTE = {
17+
"StartFunction": "#9cdcfe", # Starting function (light blue)
18+
"PyFunction": "#a277ff", # Python functions (purple)
19+
"PyClass": "#ffca85", # Python classes (orange)
20+
"ExternalModule": "#f694ff", # External module imports (pink)
21+
"HTTP_METHOD": "#ffca85", # HTTP method handlers (orange)
22+
}
23+
24+
# List of common HTTP method names to identify route handlers
25+
HTTP_METHODS = ["get", "put", "patch", "post", "head", "delete"]
26+
27+
28+
def generate_edge_meta(usage: Usage) -> dict:
29+
"""
30+
Generate metadata for graph edges based on a usage relationship.
31+
32+
Args:
33+
usage: A Usage object representing how a symbol is used
34+
35+
Returns:
36+
dict: Edge metadata including source location and symbol info
37+
"""
38+
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__}
39+
40+
41+
def is_http_method(symbol: PySymbol) -> bool:
42+
"""
43+
Check if a symbol represents an HTTP method handler.
44+
45+
Args:
46+
symbol: A Python symbol to check
47+
48+
Returns:
49+
bool: True if symbol is an HTTP method handler
50+
"""
51+
if isinstance(symbol, PyFunction) and symbol.is_method:
52+
return symbol.name in HTTP_METHODS
53+
return False
54+
55+
56+
def create_blast_radius_visualization(symbol: PySymbol, depth: int = 0):
57+
"""
58+
Recursively build a graph visualization showing how a symbol is used.
59+
Shows the "blast radius" - everything that would be affected by changes.
60+
61+
Args:
62+
symbol: Starting symbol to analyze
63+
depth: Current recursion depth
64+
"""
65+
# Stop recursion if we hit max depth
66+
if depth >= MAX_DEPTH:
67+
return
68+
69+
# Process each usage of the symbol
70+
for usage in symbol.usages:
71+
usage_symbol = usage.usage_symbol
72+
73+
# Determine node color based on symbol type
74+
if is_http_method(usage_symbol):
75+
color = COLOR_PALETTE.get("HTTP_METHOD")
76+
else:
77+
color = COLOR_PALETTE.get(usage_symbol.__class__.__name__, "#f694ff")
78+
79+
# Add node and edge to graph
80+
G.add_node(usage_symbol, color=color)
81+
G.add_edge(symbol, usage_symbol, **generate_edge_meta(usage))
82+
83+
# Recurse to process usages of this symbol
84+
create_blast_radius_visualization(usage_symbol, depth + 1)
85+
86+
87+
@codegen.function("visualize-function-blast-radius")
88+
def run(codebase: Codebase):
89+
"""
90+
Generate a visualization showing the blast radius of changes to a function.
91+
92+
This codemod:
93+
1. Identifies all usages of a target function
94+
2. Creates a graph showing how the function is used throughout the codebase
95+
3. Highlights HTTP method handlers and different types of code elements
96+
"""
97+
global G
98+
G = nx.DiGraph()
99+
100+
# Get the target function to analyze
101+
target_func = codebase.get_function("export_asset")
102+
103+
# Add starting function to graph with special color
104+
G.add_node(target_func, color=COLOR_PALETTE.get("StartFunction"))
105+
106+
# Build the visualization starting from target function
107+
create_blast_radius_visualization(target_func)
108+
109+
print(G)
110+
print("Use codegen.sh to visualize the graph!")
111+
112+
113+
if __name__ == "__main__":
114+
print("Initializing codebase...")
115+
codebase = Codebase.from_repo("codegen-oss/posthog", commit="81941c24897889a2ff2f627c693fa734967e693c", programming_language=ProgrammingLanguage.PYTHON)
116+
print(f"Codebase with {len(codebase.files)} files and {len(codebase.functions)} functions.")
117+
print("Creating graph...")
118+
119+
run(codebase)

0 commit comments

Comments
 (0)