Coverage for src/mlstatpy/graph/graphviz_helper.py: 90%

1""" 

2@file 

3@brief graphviz helper 

4""" 

5import os 

6import sys 

7from pyquickhelper.loghelper import run_cmd 

8from pyquickhelper.helpgen.conf_path_tools import find_graphviz_dot 

9 

10 

11def run_graphviz(filename, image, engine="dot"): 

12 """ 

13 Run :epkg:`GraphViz`. 

14 

15 @param filename filename which contains the graph definition 

16 @param image output image 

17 @param engine *dot* or *neato* 

18 @return output of graphviz 

19 """ 

20 ext = os.path.splitext(image)[-1] 

21 if ext != ".png": 

22 raise RuntimeError("extension should be .png not " + str(ext)) 

23 if sys.platform.startswith("win"): 

24 bin_ = os.path.dirname(find_graphviz_dot()) 

25 # if bin not in os.environ["PATH"]: 

26 # os.environ["PATH"] = os.environ["PATH"] + ";" + bin 

27 cmd = f'"{bin_}\\{engine}" -Tpng "{filename}" -o "{image}"' 

28 else: 

29 cmd = f'"{engine}" -Tpng "{filename}" -o "{image}"' 

30 out, err = run_cmd(cmd, wait=True) 

31 if len(err) > 0: 

32 raise RuntimeError( 

33 f"Unable to run Graphviz\nCMD:\n{cmd}\nOUT:\n{out}\nERR:\n{err}") 

34 return out 

35 

36 

37def edges2gv(vertices, edges): 

38 """ 

39 Converts a graph into a :epkg:`GraphViz` file format. 

40 

41 @param edges see below 

42 @param vertices see below 

43 @return gv format 

44 

45 The function creates a file ``<image>.gv``. 

46 

47 .. runpython:: 

48 :showcode: 

49 

50 from mlstatpy.graph.graphviz_helper import edges2gv 

51 gv = edges2gv([(1, "eee", "red")], 

52 [(1, 2, "blue"), (3, 4), (1, 3)]) 

53 print(gv) 

54 

55 """ 

56 memovertex = {} 

57 for v in vertices: 

58 if isinstance(v, tuple): 

59 if len(v) == 1: 

60 memovertex[v[0]] = None 

61 else: 

62 memovertex[v[0]] = v[1:] 

63 else: 

64 memovertex[v] = None 

65 for edge in edges: 

66 i, j = edge[:2] 

67 if i not in memovertex: 

68 memovertex[i] = None 

69 if j not in memovertex: 

70 memovertex[j] = None 

71 

72 li = ["digraph{"] 

73 for k, v in memovertex.items(): 

74 if v is None: 

75 li.append(f"{k} ;") 

76 elif len(v) == 1: 

77 li.append(f"\"{k}\" [label=\"{v[0]}\"];") 

78 elif len(v) == 2: 

79 li.append( 

80 f"\"{k}\" [label=\"{v[0]}\",fillcolor={v[1]},color={v[1]}];") 

81 else: 

82 raise ValueError("unable to understand " + str(v)) 

83 

84 for edge in edges: 

85 i, j = edge[:2] 

86 if len(edge) == 2: 

87 li.append(f"\"{i}\" -> \"{j}\";") 

88 elif len(edge) == 3: 

89 li.append(f"\"{i}\" -> \"{j}\" [label=\"{edge[2]}\"];") 

90 elif len(edge) == 4: 

91 li.append( 

92 f"\"{i}\" -> \"{j}\" [label=\"{edge[2]}\",color={edge[3]}];") 

93 else: 

94 raise ValueError("unable to understand " + str(edge)) 

95 li.append("}") 

96 

97 text = "\n".join(li) 

98 return text 

99 

100 

101def draw_graph_graphviz(vertices, edges, image=None, engine="dot"): 

102 """ 

103 Draws a graph using :epkg:`Graphviz`. 

104 

105 @param edges see below 

106 @param vertices see below 

107 @param image output image, None, just returns the output 

108 @param engine *dot* or *neato* 

109 @return :epkg:`Graphviz` output or 

110 the dot text if *image* is None 

111 

112 The function creates a file ``<image>.gv`` if *image* is not None. 

113 :: 

114 

115 edges = [ (1,2, label, color), (3,4), (1,3), ... ] , liste d'arcs 

116 vertices = [ (1, label, color), (2), ... ] , liste de noeuds 

117 image = nom d'image (format png) 

118 

119 """ 

120 text = edges2gv(vertices, edges) 

121 if image is None: 

122 return text 

123 filename = image + ".gv" 

124 with open(filename, "w", encoding="utf-8") as f: 

125 f.write(text) 

126 

127 out = run_graphviz(filename, image, engine=engine) 

128 if not os.path.exists(image): 

129 raise FileNotFoundError( 

130 f"GraphViz failed with no reason. '{image}' not found.") 

131 return out