Coverage for graphqler / graph / utils.py: 100%

1import networkx 

2import matplotlib.pyplot as plt 

3from matplotlib.patches import FancyBboxPatch, Patch 

4from pathlib import Path 

5 

6 

7# Switch to Agg backend to avoid GUI dependencies (e.g. Tkinter) 

8plt.switch_backend('agg') 

9 

10 

11def draw_graph(graph: networkx.DiGraph, save_path: Path): 

12 """Draws a graph with nodes as rounded rectangles and labels inside them, 

13 using professional colors suitable for academic publications, and adds a legend describing the colors. 

14 

15 Args: 

16 graph (networkx.DiGraph): The networkx graph 

17 save_path (Path): The path to save the visualization 

18 """ 

19 pos = networkx.spring_layout(graph, k=2, iterations=20) 

20 

21 # Define a professional, colorblind-friendly palette 

22 color_map = {} 

23 type_color = {} # Map from node type to color 

24 for node in graph.nodes(data=True): 

25 node_type = node[0].graphql_type 

26 if node_type == 'Mutation': 

27 color = '#1f77b4' # Blue 

28 elif node_type == 'Query': 

29 color = '#2ca02c' # Green 

30 else: 

31 color = '#7f7f7f' # Gray 

32 node_type = 'Object' # To group all other types under 'Object' 

33 color_map[node[0]] = color 

34 type_color[node_type] = color # Map node type to color 

35 

36 fig, ax = plt.subplots(figsize=(12, 8)) 

37 

38 # Detect bidirectional edges and draw them accordingly 

39 drawn_edges = set() 

40 for edge in graph.edges(): 

41 u, v = edge 

42 if (v, u) in graph.edges() and (v, u) not in drawn_edges: 

43 # Draw a single bidirectional edge 

44 x1, y1 = pos[u] 

45 x2, y2 = pos[v] 

46 ax.annotate("", 

47 xy=(x2, y2), xycoords='data', 

48 xytext=(x1, y1), textcoords='data', 

49 arrowprops=dict(arrowstyle="<->", color="black", shrinkA=15, shrinkB=15, 

50 connectionstyle="arc3,rad=0.1", linewidth=1)) 

51 drawn_edges.add((u, v)) 

52 drawn_edges.add((v, u)) 

53 elif (u, v) not in drawn_edges: 

54 # Draw a unidirectional edge 

55 x1, y1 = pos[u] 

56 x2, y2 = pos[v] 

57 ax.annotate("", 

58 xy=(x2, y2), xycoords='data', 

59 xytext=(x1, y1), textcoords='data', 

60 arrowprops=dict(arrowstyle="->", color="black", shrinkA=15, shrinkB=15, 

61 connectionstyle="arc3,rad=0.1", linewidth=1)) 

62 drawn_edges.add((u, v)) 

63 

64 # Calculate all x and y positions 

65 all_x = [pos[node][0] for node in graph.nodes()] 

66 all_y = [pos[node][1] for node in graph.nodes()] 

67 

68 # Calculate margins 

69 x_margin = (max(all_x) - min(all_x)) * 0.1 # 10% margin 

70 y_margin = (max(all_y) - min(all_y)) * 0.1 # 10% margin 

71 

72 # Set axis limits with margins 

73 ax.set_xlim(min(all_x) - x_margin, max(all_x) + x_margin) 

74 ax.set_ylim(min(all_y) - y_margin, max(all_y) + y_margin) 

75 

76 plt.axis('off') 

77 plt.tight_layout() 

78 

79 # Draw nodes as rounded rectangles with variable width 

80 # Get renderer to compute text size 

81 fig.canvas.draw() # Need to draw the figure to get the renderer 

82 renderer = fig.canvas.renderer # type: ignore 

83 

84 for node in graph.nodes(): 

85 x, y = pos[node] 

86 text = node.name 

87 

88 # Create a dummy text object to get text size 

89 text_obj = ax.text(0, 0, text, fontsize=8) 

90 bbox = text_obj.get_window_extent(renderer=renderer) 

91 # Remove the dummy text object 

92 text_obj.remove() 

93 

94 # Convert bbox width from pixels to data units 

95 inv = ax.transData.inverted() 

96 bbox_data = inv.transform([[0, 0], [bbox.width, bbox.height]]) 

97 width_data = bbox_data[1][0] - bbox_data[0][0] 

98 height_data = bbox_data[1][1] - bbox_data[0][1] 

99 

100 # Add some padding 

101 width = width_data + 0.02 * (ax.get_xlim()[1] - ax.get_xlim()[0]) 

102 height = height_data + 0.02 * (ax.get_ylim()[1] - ax.get_ylim()[0]) 

103 

104 # Center the box around (x, y) 

105 box = FancyBboxPatch((x - width / 2, y - height / 2), 

106 width, 

107 height, 

108 boxstyle="round,pad=0.02", 

109 fc=color_map[node], 

110 ec="black", 

111 linewidth=1) 

112 ax.add_patch(box) 

113 # Add label inside the box 

114 ax.text(x, y, text, horizontalalignment='center', verticalalignment='center', fontsize=8) 

115 

116 # Create custom legend handles 

117 legend_handles = [] 

118 for node_type, color in type_color.items(): 

119 patch = Patch(facecolor=color, edgecolor='black', label=node_type) 

120 legend_handles.append(patch) 

121 

122 # Add the legend to the plot 

123 ax.legend(handles=legend_handles, loc='upper right', title='Node Types', fontsize=8, title_fontsize=9) 

124 

125 # Save the figure 

126 plt.savefig(save_path, format="png", dpi=1000, bbox_inches='tight') 

127 plt.savefig(save_path.with_suffix('.pdf'), format='pdf', bbox_inches='tight') 

128 plt.close()