Codebase Visualization

This guide will show you how to create codebase visualizations using codegen.

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
• Function Dependency Graph
• Blast Radius Visualization

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)

Adding Edge Metadata

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

View on codegen.sh

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

We’ll use the same basic setup as the Call Trace Visualization tutorial.

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

View on codegen.sh

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

We’ll use the same basic setup as the Call Trace Visualization tutorial.

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

View on codegen.sh

What’s Next?

Codebase Modularity

Learn how to use Codegen to create modular codebases.

Deleting Dead Code

Learn how to use Codegen to delete dead code.

Increase Type Coverage

Learn how to use Codegen to increase type coverage.

API Reference

Explore the complete API documentation for all Codegen classes and methods.