graph_converter.py

"""
Graph Converter Module

Converts Pyan call graphs to NetworkX directed graphs for further processing.
"""

import networkx as nx
from pyan_insperation.analyzer import CallGraphVisitor


def pyan_to_networkx(pyan_graph: CallGraphVisitor) -> nx.DiGraph:
    """
    Convert a Pyan call graph to a NetworkX directed graph.
    
    This function processes both defines_edges (containment relationships like
    class-contains-method) and uses_edges (usage relationships like function-calls-function)
    from the Pyan analyzer and creates a unified NetworkX graph with labeled edges.
    
    Args:
        pyan_graph: CallGraphVisitor instance with defines_edges and uses_edges
        
    Returns:
        nx.DiGraph with nodes and labeled edges ("contains" or "use")
    """
    graph = nx.DiGraph()
    
    # Process defines_edges - containment relationships
    for node in pyan_graph.defines_edges.keys():
        graph.add_node(node)
        for defined_node in pyan_graph.defines_edges[node]:
            graph.add_node(defined_node)
            graph.add_edge(node, defined_node, label="contains")
    
    # Process uses_edges - usage relationships
    for node in pyan_graph.uses_edges.keys():
        graph.add_node(node)
        for used_node in pyan_graph.uses_edges[node]:
            graph.add_node(used_node)
            graph.add_edge(node, used_node, label="use")
    
    return graph