Visualize a graph of code cohesion

code_cohesion_visualizer.py

"""
# Create statement list

For each statment:
    Create a node

For each statement
    For each name in that statement:
        For each reference of that name to other statements:
            create an edge from statement A to statement B
"""

import black
from jedi import Script
from jedi.api.classes import Name
from pyvis.network import Network
from tqdm import tqdm

FILE_PATH = r"/Users/chrisanderson/PycharmProjects/discharge_disposition/src/models/predict_length_of_stay.py"

banned_names = ["import", "from ", "def "]

DISTINCT_COLORS = [
    "#FFB300",  # Vivid Yellow
    "#FF6800",  # Vivid Orange
    "#A6BDD7",  # Very Light Blue
    "#C10020",  # Vivid Red
    "#CEA262",  # Grayish Yellow
    "#817066",  # Medium Gray
    "#007D34",  # Vivid Green  # The following don't work well for people with defective color vision
    "#F6768E",  # Strong Purplish Pink
    "#00538A",  # Strong Blue
    "#FF7A5C",  # Strong Yellowish Pink
    # "#53377A",  # Strong Violet
    # "#FF8E00",  # Vivid Orange Yellow
    # "#B32851",  # Strong Purplish Red
    "#F4C800",  # Vivid Greenish Yellow
    # "#7F180D",  # Strong Reddish Brown
    "#93AA00",  # Vivid Yellowish Green
    "#593315",  # Deep Yellowish Brown
    "#F13A13",  # Vivid Reddish Orange
    # "#232C16",  # Dark Olive Green
]

def main():
    parent_to_color = {}
    color_index = 0

    script = Script(path=FILE_PATH)
    formatted_code = black.format_str(script._code, mode=black.FileMode(line_length=10000))
    script = Script(code=formatted_code)

    all_defined_names = script.get_names(all_scopes=True, definitions=True)
    all_referenced_names = script.get_names(all_scopes=True, references=True)

    set_of_defined_names = {x.full_name for x in all_defined_names}
    locally_defined_references = [x for x in all_referenced_names if x.full_name in set_of_defined_names]

    graph = Network()
    source_name: Name
    for source_name in tqdm(all_defined_names + locally_defined_references):
        source_node_name = name_to_node_name(source_name)

        if any(x in source_node_name for x in banned_names):
            continue

        sink_name: Name
        for sink_name in script.get_references(source_name.line, source_name.column):
            sink_node_name = name_to_node_name(sink_name)

            if any(x in sink_node_name for x in banned_names):
                continue

            source_parent_name = source_name.parent().name
            if source_parent_name in parent_to_color:
                source_color = parent_to_color[source_parent_name]
            else:
                source_color = DISTINCT_COLORS[color_index % len(DISTINCT_COLORS)]
                color_index += 1
                parent_to_color[source_parent_name] = source_color

            sink_parent_name = sink_name.parent().name
            if sink_parent_name in parent_to_color:
                sink_color = parent_to_color[sink_parent_name]
            else:
                sink_color = DISTINCT_COLORS[color_index % len(DISTINCT_COLORS)]
                color_index += 1
                parent_to_color[sink_parent_name] = sink_color

            graph.add_node(source_node_name, shape="box", color=source_color)
            graph.add_node(sink_node_name, shape="box", color=sink_color)

            if source_node_name != sink_node_name:
                graph.add_edge(source_node_name,
                               sink_node_name,
                               label=source_name.name,
                               value=abs(source_name.line - sink_name.line))

    graph.toggle_physics(True)
    graph.show_buttons(filter_=["physics"])
    graph.set_options(
        """
var options = {
  "physics": {
    "barnesHut": {
      "gravitationalConstant": -10000,
      "springLength": 700
    },
    "maxVelocity": 150,
    "minVelocity": 0.75
  },
  "configure": {}
}
"""
    )
    graph.width = "1800px"
    graph.height = "1000px"
    graph.show(
        "test.html",
    )


def name_to_node_name(name: Name):
    return f"{name.line}:{name.get_line_code().strip()}"


if __name__ == "__main__":
    main()