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>"""