graph_visualization.py

import json
from pyvis.network import Network

def create_graph(nodes, edges, physics_enabled=True):
    net = Network(notebook=True, height='100vh', width='100vw', bgcolor='#222222', font_color='white', cdn_resources='remote')

    for node in nodes:
        net.add_node(
            node['id'], 
            label=node['label'], 
            title=node['label'],
            color='blue' if node['label'] == 'OOP' else 'green'
        )

    for edge in edges:
        net.add_edge(edge['source'], edge['target'], title=edge['type'])  

    net.force_atlas_2based(
        gravity=-50,
        central_gravity=0.01,
        spring_length=100,
        spring_strength=0.08,
        damping=0.4
    )

    options = {
        "nodes": {
            "physics": physics_enabled
        },
        "edges": {
            "smooth": True
        },
        "interaction": {
            "hover": True,
            "zoomView": True
        },
        "physics": {
            "enabled": physics_enabled,
            "stabilization": {
                "enabled": True,
                "iterations": 200
            }
        }
    }

    net.set_options(json.dumps(options))
    return net

def visualize_graph(json_data, physics_enabled=True):
    if isinstance(json_data, str):
        data = json.loads(json_data)
    else:
        data = json_data
    nodes = data['nodes']
    edges = data['edges']
    net = create_graph(nodes, edges, physics_enabled)
    html = net.generate_html()
    html = html.replace("'", "\"")
    html = html.replace('<div id="mynetwork"', '<div id="mynetwork" style="height: 100vh; width: 100%;"')
    return f"""<iframe style="width: 100%; height: 100vh; border: none; margin: 0; padding: 0;" srcdoc='{html}'></iframe>"""