Skip to content

Commit d9afd02

Browse files
Add visualization code files to analyzers/visualization directory
1 parent 0cb601b commit d9afd02

File tree

9 files changed

+1348
-14
lines changed

9 files changed

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

0 commit comments

Comments
 (0)