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