Skip to content

Commit 1b9a981

Browse files
jemeza-codegencodegen-bot
and
codegen-bot
authored
Codebase visualization tutorial (#175)
# Motivation Adds important documentation on codebase visualizations # Content # Testing <!-- How was the change tested? --> # Please check the following before marking your PR as ready for review - [x] I have added tests for my changes - [x] I have updated the documentation or added new documentation as needed - [x] I have read and agree to the [Contributor License Agreement](../CLA.md) Co-authored-by: codegen-bot <[email protected]>
1 parent bd3c58c commit 1b9a981

File tree

3 files changed

+384
-0
lines changed

3 files changed

+384
-0
lines changed

docs/images/blast-radius.png

287 KB
Loading

docs/mint.json

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,7 @@
7373
"pages": [
7474
"tutorials/at-a-glance",
7575
"tutorials/migrating-apis",
76+
"tutorials/codebase-visualization",
7677
"tutorials/organize-your-codebase",
7778
"tutorials/modularity",
7879
"tutorials/deleting-dead-code",
Lines changed: 383 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,383 @@
1+
---
2+
title: "Codebase Visualization"
3+
sidebarTitle: "Codebase Visualization"
4+
description: "This guide will show you how to create codebase visualizations using [codegen](/introduction/overview)."
5+
icon: "flashlight"
6+
iconType: "solid"
7+
---
8+
9+
<Frame caption="Blast radius visualization of the `export_asset` function. Click and drag to pan, scroll to zoom.">
10+
<iframe
11+
width="100%"
12+
height="600px"
13+
scrolling="no"
14+
src={`https://codegen.sh/embedded/graph?id=347d349e-263b-481a-9601-1cd205b332b9&zoom=1&targetNodeName=export_asset`}
15+
className="rounded-xl "
16+
style={{
17+
backgroundColor: "#15141b",
18+
}}
19+
></iframe>
20+
</Frame>
21+
22+
## Overview
23+
24+
To demonstrate the visualization capabilities of the codegen we will generate three different visualizations of PostHog's open source [repository](https://github.com/PostHog/posthog).
25+
- [Call Trace Visualization](#call-trace-visualization)
26+
- [Function Dependency Graph](#function-dependency-graph)
27+
- [Blast Radius Visualization](#blast-radius-visualization)
28+
29+
30+
## Call Trace Visualization
31+
32+
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](https://github.com/PostHog/posthog/blob/c2986d9ac7502aa107a4afbe31b3633848be6582/posthog/api/sharing.py#L163).
33+
34+
35+
### Basic Setup
36+
First, we'll set up our codebase, graph and configure some basic parameters:
37+
38+
```python
39+
import networkx as nx
40+
from codegen import Codebase
41+
42+
# Initialize codebase
43+
codebase = Codebase("path/to/posthog/")
44+
45+
# Create a directed graph for representing call relationships
46+
G = nx.DiGraph()
47+
48+
# Configuration flags
49+
IGNORE_EXTERNAL_MODULE_CALLS = True # Skip calls to external modules
50+
IGNORE_CLASS_CALLS = False # Include class definition calls
51+
MAX_DEPTH = 10
52+
53+
COLOR_PALETTE = {
54+
"StartFunction": "#9cdcfe", # Light blue - Start Function
55+
"PyFunction": "#a277ff", # Soft purple/periwinkle - PyFunction
56+
"PyClass": "#ffca85", # Warm peach/orange - PyClass
57+
"ExternalModule": "#f694ff" # Bright magenta/pink - ExternalModule
58+
}
59+
```
60+
61+
### Building the Visualization
62+
We'll create a function that will recursively traverse the call trace of a function and add nodes and edges to the graph:
63+
64+
```python
65+
def create_downstream_call_trace(src_func: Function, depth: int = 0):
66+
"""Creates call graph by recursively traversing function calls
67+
68+
Args:
69+
src_func (Function): Starting function for call graph
70+
depth (int): Current recursion depth
71+
"""
72+
# Prevent infinite recursion
73+
if MAX_DEPTH <= depth:
74+
return
75+
76+
# External modules are not functions
77+
if isinstance(src_func, ExternalModule):
78+
return
79+
80+
# Process each function call
81+
for call in src_func.function_calls:
82+
# Skip self-recursive calls
83+
if call.name == src_func.name:
84+
continue
85+
86+
# Get called function definition
87+
func = call.function_definition
88+
if not func:
89+
continue
90+
91+
# Apply configured filters
92+
if isinstance(func, ExternalModule) and IGNORE_EXTERNAL_MODULE_CALLS:
93+
continue
94+
if isinstance(func, Class) and IGNORE_CLASS_CALLS:
95+
continue
96+
97+
# Generate display name (include class for methods)
98+
if isinstance(func, Class) or isinstance(func, ExternalModule):
99+
func_name = func.name
100+
elif isinstance(func, Function):
101+
func_name = f"{func.parent_class.name}.{func.name}" if func.is_method else func.name
102+
103+
# Add node and edge with metadata
104+
G.add_node(func, name=func_name,
105+
color=COLOR_PALETTE.get(func.__class__.__name__))
106+
G.add_edge(src_func, func, **generate_edge_meta(call))
107+
108+
# Recurse for regular functions
109+
if isinstance(func, Function):
110+
create_downstream_call_trace(func, depth + 1)
111+
```
112+
113+
### Adding Edge Metadata
114+
We can enrich our edges with metadata about the function calls:
115+
116+
```python
117+
def generate_edge_meta(call: FunctionCall) -> dict:
118+
"""Generate metadata for call graph edges
119+
120+
Args:
121+
call (FunctionCall): Function call information
122+
123+
Returns:
124+
dict: Edge metadata including name and location
125+
"""
126+
return {
127+
"name": call.name,
128+
"file_path": call.filepath,
129+
"start_point": call.start_point,
130+
"end_point": call.end_point,
131+
"symbol_name": "FunctionCall"
132+
}
133+
```
134+
### Visualizing the Graph
135+
Finally, we can visualize our call graph starting from a specific function:
136+
```python
137+
# Get target function to analyze
138+
target_class = codebase.get_class('SharingConfigurationViewSet')
139+
target_method = target_class.get_method('patch')
140+
141+
# Add root node
142+
G.add_node(target_method,
143+
name=f"{target_class.name}.{target_method.name}",
144+
color=COLOR_PALETTE["StartFunction"])
145+
146+
# Build the call graph
147+
create_downstream_call_trace(target_method)
148+
149+
# Render the visualization
150+
codebase.visualize(G)
151+
```
152+
153+
### Common Use Cases
154+
The call graph visualization is particularly useful for:
155+
- Understanding complex codebases
156+
- Planning refactoring efforts
157+
- Identifying tightly coupled components
158+
- Analyzing critical paths
159+
- Documenting system architecture
160+
161+
### Take a look
162+
<iframe
163+
width="100%"
164+
height="600px"
165+
scrolling="no"
166+
src={`https://codegen.sh/embedded/graph?id=9a141c52-c0c2-4737-bdc3-5d164791d3b5&zoom=1&targetNodeName=SharingConfigurationViewSet.patch`}
167+
className="rounded-xl "
168+
style={{
169+
backgroundColor: "#15141b",
170+
}}
171+
></iframe>
172+
173+
## Function Dependency Graph
174+
175+
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](https://github.com/PostHog/posthog/blob/c2986d9ac7502aa107a4afbe31b3633848be6582/posthog/hogql_queries/query_runner.py#L152).
176+
177+
### Basic Setup
178+
<Info>
179+
We'll use the same basic setup as the [Call Trace Visualization](/tutorials/codebase-visualization#call-trace-visualization) tutorial.
180+
</Info>
181+
182+
### Building the Dependency Graph
183+
The core function for building our dependency graph:
184+
```python
185+
def create_dependencies_visualization(symbol: Symbol, depth: int = 0):
186+
"""Creates visualization of symbol dependencies
187+
188+
Args:
189+
symbol (Symbol): Starting symbol to analyze
190+
depth (int): Current recursion depth
191+
"""
192+
# Prevent excessive recursion
193+
if depth >= MAX_DEPTH:
194+
return
195+
196+
# Process each dependency
197+
for dep in symbol.dependencies:
198+
dep_symbol = None
199+
200+
# Handle different dependency types
201+
if isinstance(dep, Symbol):
202+
# Direct symbol reference
203+
dep_symbol = dep
204+
elif isinstance(dep, Import):
205+
# Import statement - get resolved symbol
206+
dep_symbol = dep.resolved_symbol if dep.resolved_symbol else None
207+
208+
if dep_symbol:
209+
# Add node with appropriate styling
210+
G.add_node(dep_symbol,
211+
color=COLOR_PALETTE.get(dep_symbol.__class__.__name__,
212+
"#f694ff"))
213+
214+
# Add dependency relationship
215+
G.add_edge(symbol, dep_symbol)
216+
217+
# Recurse unless it's a class (avoid complexity)
218+
if not isinstance(dep_symbol, PyClass):
219+
create_dependencies_visualization(dep_symbol, depth + 1)
220+
```
221+
222+
### Visualizing the Graph
223+
Finally, we can visualize our dependency graph starting from a specific symbol:
224+
```python
225+
# Get target symbol
226+
target_func = codebase.get_function("get_query_runner")
227+
228+
# Add root node
229+
G.add_node(target_func, color=COLOR_PALETTE["StartFunction"])
230+
231+
# Generate dependency graph
232+
create_dependencies_visualization(target_func)
233+
234+
# Render visualization
235+
codebase.visualize(G)
236+
```
237+
238+
### Take a look
239+
<iframe
240+
width="100%"
241+
height="600px"
242+
scrolling="no"
243+
src={`https://codegen.sh/embedded/graph?id=bb7b227b-cc89-4e92-b71f-fb0d6265eb3d&zoom=0.8&targetNodeName=get_query_runner`}
244+
className="rounded-xl "
245+
style={{
246+
backgroundColor: "#15141b",
247+
}}
248+
></iframe>
249+
250+
## Blast Radius visualization
251+
252+
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](https://github.com/PostHog/posthog/blob/c2986d9ac7502aa107a4afbe31b3633848be6582/posthog/tasks/exporter.py#L57).
253+
254+
### Basic Setup
255+
<Info>
256+
We'll use the same basic setup as the [Call Trace Visualization](/tutorials/codebase-visualization#call-trace-visualization) tutorial.
257+
</Info>
258+
259+
### Helper Functions
260+
We'll create some utility functions to help build our visualization:
261+
```python
262+
# List of HTTP methods to highlight
263+
HTTP_METHODS = ["get", "put", "patch", "post", "head", "delete"]
264+
265+
def generate_edge_meta(usage: Usage) -> dict:
266+
"""Generate metadata for graph edges
267+
268+
Args:
269+
usage (Usage): Usage relationship information
270+
271+
Returns:
272+
dict: Edge metadata including name and location
273+
"""
274+
return {
275+
"name": usage.match.source,
276+
"file_path": usage.match.filepath,
277+
"start_point": usage.match.start_point,
278+
"end_point": usage.match.end_point,
279+
"symbol_name": usage.match.__class__.__name__
280+
}
281+
282+
def is_http_method(symbol: PySymbol) -> bool:
283+
"""Check if a symbol is an HTTP endpoint method
284+
285+
Args:
286+
symbol (PySymbol): Symbol to check
287+
288+
Returns:
289+
bool: True if symbol is an HTTP method
290+
"""
291+
if isinstance(symbol, PyFunction) and symbol.is_method:
292+
return symbol.name in HTTP_METHODS
293+
return False
294+
```
295+
296+
### Building the Blast Radius Visualization
297+
The main function for creating our blast radius visualization:
298+
```python
299+
def create_blast_radius_visualization(symbol: PySymbol, depth: int = 0):
300+
"""Create visualization of symbol usage relationships
301+
302+
Args:
303+
symbol (PySymbol): Starting symbol to analyze
304+
depth (int): Current recursion depth
305+
"""
306+
# Prevent excessive recursion
307+
if depth >= MAX_DEPTH:
308+
return
309+
310+
# Process each usage of the symbol
311+
for usage in symbol.usages:
312+
usage_symbol = usage.usage_symbol
313+
314+
# Determine node color based on type
315+
if is_http_method(usage_symbol):
316+
color = COLOR_PALETTE.get("HTTP_METHOD")
317+
else:
318+
color = COLOR_PALETTE.get(usage_symbol.__class__.__name__, "#f694ff")
319+
320+
# Add node and edge to graph
321+
G.add_node(usage_symbol, color=color)
322+
G.add_edge(symbol, usage_symbol, **generate_edge_meta(usage))
323+
324+
# Recursively process usage symbol
325+
create_blast_radius_visualization(usage_symbol, depth + 1)
326+
```
327+
328+
### Visualizing the Graph
329+
Finally, we can create our blast radius visualization:
330+
```python
331+
# Get target function to analyze
332+
target_func = codebase.get_function('export_asset')
333+
334+
# Add root node
335+
G.add_node(target_func, color=COLOR_PALETTE.get("StartFunction"))
336+
337+
# Build the visualization
338+
create_blast_radius_visualization(target_func)
339+
340+
# Render reversed graph to show impact flow
341+
codebase.visualize(G.reverse())
342+
```
343+
344+
### Take a look
345+
<iframe
346+
width="100%"
347+
height="600px"
348+
scrolling="no"
349+
src={`https://codegen.sh/embedded/graph?id=347d349e-263b-481a-9601-1cd205b332b9&zoom=1&targetNodeName=export_asset`}
350+
className="rounded-xl "
351+
style={{
352+
backgroundColor: "#15141b",
353+
}}
354+
></iframe>
355+
356+
## What's Next?
357+
358+
<CardGroup cols={2}>
359+
<Card
360+
title="Codebase Modularity"
361+
icon="diagram-project"
362+
href="/tutorials/modularity"
363+
>
364+
Learn how to use Codegen to create modular codebases.
365+
</Card>
366+
<Card
367+
title="Deleting Dead Code"
368+
icon="trash"
369+
href="/tutorials/deleting-dead-code"
370+
>
371+
Learn how to use Codegen to delete dead code.
372+
</Card>
373+
<Card
374+
title="Increase Type Coverage"
375+
icon="shield-check"
376+
href="/tutorials/increase-type-coverage"
377+
>
378+
Learn how to use Codegen to increase type coverage.
379+
</Card>
380+
<Card title="API Reference" icon="code" href="/api-reference">
381+
Explore the complete API documentation for all Codegen classes and methods.
382+
</Card>
383+
</CardGroup>

0 commit comments

Comments
 (0)