forked from codegen-sh/codegen
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathvisualization_manager.py
More file actions
63 lines (49 loc) · 1.88 KB
/
Copy pathvisualization_manager.py
File metadata and controls
63 lines (49 loc) · 1.88 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
import os
import plotly.graph_objects as go
from networkx import Graph
from codegen.git.repo_operator.repo_operator import RepoOperator
from codegen.sdk.core.interfaces.editable import Editable
from codegen.shared.logging.get_logger import get_logger
from codegen.visualizations.viz_utils import graph_to_json
logger = get_logger(__name__)
class VisualizationManager:
op: RepoOperator
def __init__(
self,
op: RepoOperator,
) -> None:
self.op = op
@property
def viz_path(self) -> str:
return os.path.join(self.op.base_dir, "codegen-graphviz")
@property
def viz_file_path(self) -> str:
return os.path.join(self.viz_path, "graph.json")
def clear_graphviz_data(self) -> None:
if self.op.folder_exists(self.viz_path):
self.op.emptydir(self.viz_path)
def write_graphviz_data(self, G: Graph | go.Figure, root: Editable | str | int | None = None) -> None:
"""Writes the graph data to a file.
Args:
----
G (Graph | go.Figure): A NetworkX Graph object representing the graph to be visualized.
root (str | None): The root node to visualize. Defaults to None.
Returns:
------
None
"""
# Convert the graph to a JSON-serializable format
if isinstance(G, Graph):
graph_json = graph_to_json(G, root)
elif isinstance(G, go.Figure):
graph_json = G.to_json()
# Check if the visualization path exists, if so, empty it
if self.op.folder_exists(self.viz_path):
self.op.emptydir(self.viz_path)
else:
# If the path doesn't exist, create it
self.op.mkdir(self.viz_path)
# Write the graph data to a file
with open(self.viz_file_path, "w") as f:
f.write(graph_json)
f.flush() # Ensure data is written to disk