Codebase Visualization

Codegen provides the ability to create interactive graph visualizations via the codebase.visualize(…) method.

These visualizations have a number of applications, including:

• Understanding codebase structure
• Monitoring critical code paths
• Analyzing dependencies
• Understanding inheritance hierarchies

This guide provides a basic overview of graph creation and customization. Like the one below which displays the call_graph for the modal/client.py module.

Codegen visualizations are powered by NetworkX and rendered using d3.

Basic Usage

The Codebase.visualize method operates on a NetworkX DiGraph.

import networkx as nx

# Basic visualization
G = nx.grid_2d_graph(5, 5)
# Or start with an empty graph
# G = nx.DiGraph()
codebase.visualize(G)

It is up to the developer to add nodes and edges to the graph.

Adding Nodes and Edges

When adding nodes to your graph, you can either add the symbol directly or just its name:

import networkx as nx
G = nx.DiGraph()
function = codebase.get_function("my_function")

# Add the function object directly - enables source code preview
graph.add_node(function)  # Will show function's source code on click

# Add just the name - no extra features
graph.add_node(function.name)  # Will only show the name

Adding symbols to the graph directly (as opposed to adding by name) enables automatic type information, code preview on hover, and more.

Common Visualization Types

Call Graphs

Visualize how functions call each other and trace execution paths:

def create_call_graph(entry_point: Function):
    graph = nx.DiGraph()

    def add_calls(func):
        for call in func.call_sites:
            called_func = call.resolved_symbol
            if called_func:
                # Add function objects for rich previews
                graph.add_node(func)
                graph.add_node(called_func)
                graph.add_edge(func, called_func)
                add_calls(called_func)

    add_calls(entry_point)
    return graph

# Visualize API endpoint call graph
endpoint = codebase.get_function("handle_request")
call_graph = create_call_graph(endpoint)
codebase.visualize(call_graph, root=endpoint)

Learn more about traversing the call graph here.

React Component Trees

Visualize the hierarchy of React components:

def create_component_tree(root_component: Class):
    graph = nx.DiGraph()

    def add_children(component):
        for usage in component.usages:
            if isinstance(usage.parent, Class) and "Component" in usage.parent.bases:
                graph.add_edge(component.name, usage.parent.name)
                add_children(usage.parent)

    add_children(root_component)
    return graph

# Visualize component hierarchy
app = codebase.get_class("App")
component_tree = create_component_tree(app)
codebase.visualize(component_tree, root=app)

Inheritance Graphs

Visualize class inheritance relationships:

import networkx as nx

G = nx.DiGraph()
base = codebase.get_class("BaseModel")

def add_subclasses(cls):
    for subclass in cls.subclasses:
        G.add_edge(cls, subclass)
        add_subclasses(subclass)

add_subclasses(base)

codebase.visualize(G, root=base)

Module Dependencies

Visualize dependencies between modules:

def create_module_graph(start_file: File):
    G = nx.DiGraph()

    def add_imports(file):
        for imp in file.imports:
            if imp.resolved_symbol and imp.resolved_symbol.file:
                graph.add_edge(file, imp.resolved_symbol.file)
                add_imports(imp.resolved_symbol.file)

    add_imports(start_file)
    return graph

# Visualize module dependencies
main = codebase.get_file("main.py")
module_graph = create_module_graph(main)
codebase.visualize(module_graph, root=main)

Function Modularity

Visualize function groupings by modularity:

def create_modularity_graph(functions: list[Function]):
    graph = nx.Graph()

    # Group functions by shared dependencies
    for func in functions:
        for dep in func.dependencies:
            if isinstance(dep, Function):
                weight = len(set(func.dependencies) & set(dep.dependencies))
                if weight > 0:
                    graph.add_edge(func.name, dep.name, weight=weight)

    return graph

# Visualize function modularity
funcs = codebase.functions
modularity_graph = create_modularity_graph(funcs)
codebase.visualize(modularity_graph)

Customizing Visualizations

You can customize your visualizations using NetworkX’s attributes while still preserving the smart node features:

def create_custom_graph(codebase):
    graph = nx.DiGraph()

    # Add nodes with custom attributes while preserving source preview
    for func in codebase.functions:
        graph.add_node(func,
            color='red' if func.is_public else 'blue',
            shape='box' if func.is_async else 'oval'
        )

    # Add edges between actual function objects
    for func in codebase.functions:
        for call in func.call_sites:
            if call.resolved_symbol:
                graph.add_edge(func, call.resolved_symbol,
                    style='dashed' if call.is_conditional else 'solid',
                    weight=call.count
                )

    return graph

Best Practices

1. Use Symbol Objects for Rich Features

   # Better: Add symbol objects for rich previews
   # This will include source code previews, syntax highlighting, type information, etc.
   for func in api_funcs:
       graph.add_node(func)
   
   # Basic: Just names, no extra features
   for func in api_funcs:
       graph.add_node(func.name)
   

2. Focus on Relevant Subgraphs

   # Better: Visualize specific subsystem
   api_funcs = [f for f in codebase.functions if "api" in f.filepath]
   api_graph = create_call_graph(api_funcs)
   codebase.visualize(api_graph)
   
   # Avoid: Visualizing entire codebase
   full_graph = create_call_graph(codebase.functions)  # Too complex
   

3. Use Meaningful Layouts

   # Group related nodes together
   graph.add_node(controller_class, cluster="api")
   graph.add_node(service_class, cluster="db")
   

4. Add Visual Hints

   # Color code by type while preserving rich previews
   for node in codebase.functions:
       if "Controller" in node.name:
           graph.add_node(node, color="red")
       elif "Service" in node.name:
           graph.add_node(node, color="blue")
   

Limitations

• Large graphs may become difficult to read
• Complex relationships might need multiple views
• Some graph layouts may take time to compute
• Preview features only work when adding symbol objects directly