src/visualization.py

from st_link_analysis import EdgeStyle, NodeStyle
from src.config import config

def get_node_styles() -> list[NodeStyle]:
    node_styles = []
    for class_name in config.ICON_MAPPING.keys():
        color = config.COLOR_MAPPING.get(class_name, "#888888")
        icon = config.ICON_MAPPING.get(class_name, None)
        node_styles.append(NodeStyle(
            label=class_name,
            color=color,
            icon=icon,
        ))
    return node_styles

def get_edge_styles() -> list[EdgeStyle]:
    edge_styles = [
        EdgeStyle(
            label="CITES",
        )
    ]

    return edge_styles